ClaudeDo/src/ClaudeDo.Data/Repositories/TaskRepository.cs

using ClaudeDo.Data.Models;
using Microsoft.EntityFrameworkCore;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Data.Repositories;

public sealed class TaskRepository
{
    private readonly ClaudeDoDbContext _context;
    private readonly AttachmentStore _attachments;

    public TaskRepository(ClaudeDoDbContext context, AttachmentStore? attachments = null)
    {
        _context = context;
        _attachments = attachments ?? new AttachmentStore();
    }

    #region CRUD

    public async Task AddAsync(TaskEntity entity, CancellationToken ct = default)
    {
        // Append at bottom of the list by default: SortOrder = max(listId) + 1.
        var maxSort = await _context.Tasks
            .Where(t => t.ListId == entity.ListId)
            .Select(t => (int?)t.SortOrder)
            .MaxAsync(ct);
        entity.SortOrder = (maxSort ?? -1) + 1;

        await TaskNumberAllocator.AddWithNumberAsync(_context, entity, ct);
    }

    public async Task UpdateAsync(TaskEntity entity, CancellationToken ct = default)
    {
        var tracked = _context.Tasks.Local.FirstOrDefault(t => t.Id == entity.Id);
        if (tracked is not null && !ReferenceEquals(tracked, entity))
            _context.Entry(tracked).State = Microsoft.EntityFrameworkCore.EntityState.Detached;
        // Entry().State instead of Update(): Update marks the whole loaded graph Modified, which
        // for the all-key task_dependencies rows means an UPDATE with nothing to set. Dependency
        // edges are written only by TaskStateService.SetDependsOnAsync.
        _context.Entry(entity).State = Microsoft.EntityFrameworkCore.EntityState.Modified;
        await _context.SaveChangesAsync(ct);
    }

    public async Task DeleteAsync(string taskId, CancellationToken ct = default)
    {
        await _context.Tasks.Where(t => t.Id == taskId).ExecuteDeleteAsync(ct);
        _attachments.DeleteTaskDir(taskId);
    }

    public async Task<TaskEntity?> GetByIdAsync(string taskId, CancellationToken ct = default)
    {
        return await _context.Tasks.AsNoTracking().Include(t => t.Dependencies)
            .FirstOrDefaultAsync(t => t.Id == taskId, ct);
    }

    public async Task<TaskEntity?> GetByNumberAsync(int number, CancellationToken ct = default)
    {
        return await _context.Tasks.AsNoTracking().Include(t => t.Dependencies)
            .FirstOrDefaultAsync(t => t.Number == number, ct);
    }

    public async Task<List<TaskEntity>> GetByListIdAsync(string listId, CancellationToken ct = default)
    {
        return await _context.Tasks
            .Include(t => t.Dependencies)
            .Where(t => t.ListId == listId)
            .OrderBy(t => t.SortOrder).ThenBy(t => t.CreatedAt)
            .ToListAsync(ct);
    }

    /// <summary>
    /// Renumbers tasks in a list to 0..N-1 according to <paramref name="orderedTaskIds"/>.
    /// Ids not belonging to the list are ignored; ids missing from the list are untouched.
    /// </summary>
    public async Task ReorderAsync(string listId, IReadOnlyList<string> orderedTaskIds, CancellationToken ct = default)
    {
        if (orderedTaskIds.Count == 0) return;

        var idSet = orderedTaskIds.ToHashSet();
        var tasks = await _context.Tasks
            .Where(t => t.ListId == listId && idSet.Contains(t.Id))
            .ToListAsync(ct);

        for (int i = 0; i < orderedTaskIds.Count; i++)
        {
            var task = tasks.FirstOrDefault(t => t.Id == orderedTaskIds[i]);
            if (task is not null) task.SortOrder = i;
        }

        await _context.SaveChangesAsync(ct);
    }

    /// <summary>
    /// Returns the ids of every descendant of <paramref name="taskId"/> (children, grandchildren, ...),
    /// walking the ParentTaskId chain breadth-first. Does not include the task itself.
    /// </summary>
    public async Task<List<string>> GetDescendantIdsAsync(string taskId, CancellationToken ct = default)
    {
        var result = new List<string>();
        // Guards against a corrupt parent chain (a self-parent or a cycle would loop forever).
        var seen = new HashSet<string>(StringComparer.Ordinal) { taskId };
        var frontier = new List<string> { taskId };
        while (frontier.Count > 0)
        {
            var children = await _context.Tasks.AsNoTracking()
                .Where(t => t.ParentTaskId != null && frontier.Contains(t.ParentTaskId))
                .Select(t => t.Id)
                .ToListAsync(ct);
            var fresh = children.Where(seen.Add).ToList();
            if (fresh.Count == 0) break;
            result.AddRange(fresh);
            frontier = fresh;
        }
        return result;
    }

    /// <summary>
    /// Moves a task (and every descendant, so a child never ends up in a different list than its
    /// parent) to <paramref name="targetListId"/>, appending it at the end of the target list's
    /// order. ListId is init-only, so the move goes through ExecuteUpdate rather than a tracked
    /// entity mutation.
    /// </summary>
    public async Task MoveToListAsync(string taskId, string targetListId, CancellationToken ct = default)
    {
        var exists = await _context.Tasks.AsNoTracking().AnyAsync(t => t.Id == taskId, ct);
        if (!exists)
            throw new InvalidOperationException($"Task {taskId} not found.");

        var descendantIds = await GetDescendantIdsAsync(taskId, ct);
        var movedIds = new List<string> { taskId };
        movedIds.AddRange(descendantIds);

        var maxSort = await _context.Tasks
            .Where(t => t.ListId == targetListId)
            .Select(t => (int?)t.SortOrder)
            .MaxAsync(ct);

        await _context.Tasks
            .Where(t => movedIds.Contains(t.Id))
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.ListId, targetListId), ct);

        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.SortOrder, (maxSort ?? -1) + 1), ct);
    }

    /// <summary>
    /// Returns all tasks that qualify as "real" Idle backlog items for online mirroring:
    /// Status==Idle, no parent, PlanningPhase==None, not blocked.
    /// </summary>
    public async Task<List<TaskEntity>> GetAllIdleBacklogAsync(CancellationToken ct = default)
    {
        return await _context.Tasks
            .AsNoTracking()
            .Where(t => t.Status == TaskStatus.Idle
                        && t.ParentTaskId == null
                        && t.PlanningPhase == PlanningPhase.None
                        && t.BlockedByTaskId == null)
            .OrderBy(t => t.SortOrder).ThenBy(t => t.CreatedAt)
            .ToListAsync(ct);
    }

    #endregion

    #region Status transitions

    internal async Task MarkRunningAsync(string taskId, DateTime startedAt, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Running)
                .SetProperty(t => t.StartedAt, startedAt), ct);
    }

    internal async Task MarkDoneAsync(string taskId, DateTime finishedAt, string? result, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Done)
                .SetProperty(t => t.FinishedAt, finishedAt)
                .SetProperty(t => t.Result, result), ct);
    }

    internal async Task MarkFailedAsync(string taskId, DateTime finishedAt, string? result, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Failed)
                .SetProperty(t => t.FinishedAt, finishedAt)
                .SetProperty(t => t.Result, result), ct);
    }

    public async Task SetLogPathAsync(string taskId, string logPath, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.LogPath, logPath), ct);
    }

    public async Task SetRoadblockCountAsync(string taskId, int count, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.RoadblockCount, count), ct);
    }

    /// Tasks eligible for UsageLimitAutoContinueCoordinator: failed on the usage limit and never
    /// auto-continued before (that flag is never reset, so a task is only offered this once).
    public async Task<List<TaskEntity>> GetUsageLimitAutoContinueCandidatesAsync(CancellationToken ct = default)
    {
        return await _context.Tasks
            .Where(t => t.Status == TaskStatus.Failed
                && t.FailureReason == "usage_limit"
                && t.UsageLimitAutoContinuedAt == null)
            .ToListAsync(ct);
    }

    /// Atomically claims the auto-continue slot for one task. Returns false if it was already
    /// claimed (e.g. by an earlier tick) — the caller must not fire ContinueTask in that case.
    public async Task<bool> TryClaimUsageLimitAutoContinueAsync(string taskId, DateTime firedAtUtc, CancellationToken ct = default)
    {
        var affected = await _context.Tasks
            .Where(t => t.Id == taskId && t.UsageLimitAutoContinuedAt == null)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.UsageLimitAutoContinuedAt, firedAtUtc), ct);
        return affected > 0;
    }

    /// Undoes a claim whose ContinueTask never actually started (the single override slot was
    /// busy, the run was rejected). Without this the task keeps a marker for a run that never
    /// happened and the candidate query excludes it forever.
    public async Task ReleaseUsageLimitAutoContinueClaimAsync(string taskId, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.UsageLimitAutoContinuedAt, (DateTime?)null), ct);
    }

    internal async Task<int> FlipAllRunningToFailedAsync(string reason, CancellationToken ct = default)
    {
        var resultText = "[stale] " + reason;
        var now = DateTime.UtcNow;
        return await _context.Tasks
            .Where(t => t.Status == TaskStatus.Running)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Failed)
                .SetProperty(t => t.FinishedAt, now)
                .SetProperty(t => t.Result, resultText), ct);
    }

    public async Task ResetToManualAsync(string taskId, CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Idle)
                .SetProperty(t => t.StartedAt, (DateTime?)null)
                .SetProperty(t => t.FinishedAt, (DateTime?)null)
                .SetProperty(t => t.Result, (string?)null), ct);
    }

    #endregion

    #region Agent settings

    public async Task UpdateAgentSettingsAsync(
        string taskId,
        string? model,
        string? systemPrompt,
        string? agentPath,
        int? maxTurns = null,
        string? sessionSkills = null,
        string? permissionMode = null,
        CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Model, model)
                .SetProperty(t => t.SystemPrompt, systemPrompt)
                .SetProperty(t => t.AgentPath, agentPath)
                .SetProperty(t => t.MaxTurns, maxTurns)
                .SetProperty(t => t.SessionSkills, sessionSkills)
                .SetProperty(t => t.PermissionMode, permissionMode), ct);
    }

    #endregion

    #region Planning

    public async Task<List<TaskEntity>> GetChildrenAsync(string parentId, CancellationToken ct = default)
    {
        return await _context.Tasks
            .AsNoTracking()
            .Where(t => t.ParentTaskId == parentId)
            .OrderBy(t => t.SortOrder).ThenBy(t => t.CreatedAt)
            .ToListAsync(ct);
    }

    public async Task<TaskEntity> CreateChildAsync(
        string parentId,
        string title,
        string? description,
        string? commitType,
        string? createdBy = null,
        string? model = null,
        int? maxTurns = null,
        CancellationToken ct = default)
    {
        if (maxTurns.HasValue && maxTurns.Value <= 0)
            throw new ArgumentException($"maxTurns must be positive, got {maxTurns.Value}.");

        // AsNoTracking: SetPlanningStartedAsync mutates via ExecuteUpdate which
        // bypasses the change tracker; a tracked Find would return stale data.
        var parent = await _context.Tasks.AsNoTracking()
            .FirstOrDefaultAsync(t => t.Id == parentId, ct);
        if (parent is null)
            throw new InvalidOperationException($"Parent task {parentId} not found.");

        var maxSort = await _context.Tasks
            .Where(t => t.ListId == parent.ListId)
            .Select(t => (int?)t.SortOrder)
            .MaxAsync(ct);

        var child = new TaskEntity
        {
            Id = Guid.NewGuid().ToString(),
            ListId = parent.ListId,
            Title = title,
            Description = description,
            Status = TaskStatus.Idle,
            CreatedAt = DateTime.UtcNow,
            CommitType = string.IsNullOrEmpty(commitType) ? parent.CommitType : commitType,
            ParentTaskId = parentId,
            SortOrder = (maxSort ?? -1) + 1,
            CreatedBy = createdBy,
            Model = ModelRegistry.NormalizeAlias(model),
            MaxTurns = maxTurns,
        };
        await TaskNumberAllocator.AddWithNumberAsync(_context, child, ct);
        return child;
    }

    public async Task UpdateChildAsync(
        string taskId,
        string? title,
        string? description,
        string? commitType,
        TaskStatus? status,
        CancellationToken ct = default)
    {
        var task = await _context.Tasks.FirstOrDefaultAsync(t => t.Id == taskId, ct)
            ?? throw new InvalidOperationException($"Task {taskId} not found.");

        if (title is not null) task.Title = title;
        if (description is not null) task.Description = description;
        if (commitType is not null) task.CommitType = commitType;
        if (status.HasValue) task.Status = status.Value;

        await _context.SaveChangesAsync(ct);
    }

    public async Task UpdatePlanningTaskAsync(
        string taskId,
        string? title,
        string? description,
        CancellationToken ct = default)
    {
        var entity = await _context.Tasks.AsNoTracking().FirstOrDefaultAsync(t => t.Id == taskId, ct)
            ?? throw new InvalidOperationException("Planning task not found.");
        if (title is not null) entity.Title = title;
        if (description is not null) entity.Description = description;
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Title, entity.Title)
                .SetProperty(t => t.Description, entity.Description), ct);
    }

    public async Task<TaskEntity?> SetPlanningStartedAsync(
        string taskId,
        string sessionToken,
        CancellationToken ct = default)
    {
        var affected = await _context.Tasks
            .Where(t => t.Id == taskId
                        && t.Status == TaskStatus.Idle
                        && t.PlanningPhase == PlanningPhase.None)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.PlanningPhase, PlanningPhase.Active)
                .SetProperty(t => t.PlanningSessionToken, sessionToken), ct);

        if (affected == 0) return null;
        return await _context.Tasks.AsNoTracking().FirstOrDefaultAsync(t => t.Id == taskId, ct);
    }

    public async Task SetPlanningSessionTokenAsync(
        string taskId,
        string sessionToken,
        CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.PlanningSessionToken, sessionToken), ct);
    }

    public async Task UpdatePlanningSessionIdAsync(
        string parentId,
        string sessionId,
        CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == parentId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.PlanningSessionId, sessionId), ct);
    }

    // Stamps the review range's head commit for a worktree-less "list handler" host task
    // (Mission Control's "Let Claude handle it") when its ConPTY session is submitted for review.
    public async Task SetHandlerHeadCommitAsync(
        string taskId,
        string headCommit,
        CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.HandlerHeadCommit, headCommit), ct);
    }

    // Persists the claude session id a fresh embedded ConPTY interactive task session will run
    // under, written BEFORE launch so a closed/aborted session still leaves a resumable id.
    public async Task SetInteractiveSessionIdAsync(
        string taskId,
        string? sessionId,
        CancellationToken ct = default)
    {
        await _context.Tasks
            .Where(t => t.Id == taskId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.InteractiveSessionId, sessionId), ct);
    }

    public async Task<TaskEntity?> FindByPlanningTokenAsync(
        string token,
        CancellationToken ct = default)
    {
        if (string.IsNullOrEmpty(token)) return null;
        return await _context.Tasks
            .AsNoTracking()
            .FirstOrDefaultAsync(t => t.PlanningSessionToken == token, ct);
    }

    public async Task<DiscardPlanningOutcome> DiscardPlanningAsync(
        string parentId,
        bool dequeueQueuedChildren,
        CancellationToken ct = default)
    {
        using var tx = await _context.Database.BeginTransactionAsync(ct);

        var parent = await _context.Tasks
            .AsNoTracking()
            .FirstOrDefaultAsync(t => t.Id == parentId, ct);
        if (parent is null || parent.PlanningPhase != PlanningPhase.Active)
        {
            await tx.RollbackAsync(ct);
            return new DiscardPlanningOutcome(DiscardPlanningResult.NotInPlanning, 0, 0);
        }

        var children = await _context.Tasks
            .Where(t => t.ParentTaskId == parentId)
            .Select(t => new { t.Id, t.Status })
            .ToListAsync(ct);

        var runningCount = children.Count(c => c.Status == TaskStatus.Running);
        if (runningCount > 0)
        {
            await tx.RollbackAsync(ct);
            return new DiscardPlanningOutcome(DiscardPlanningResult.BlockedByRunningChildren, 0, runningCount);
        }

        var queuedIds = children.Where(c => c.Status == TaskStatus.Queued).Select(c => c.Id).ToList();
        if (queuedIds.Count > 0)
        {
            if (!dequeueQueuedChildren)
            {
                await tx.RollbackAsync(ct);
                return new DiscardPlanningOutcome(DiscardPlanningResult.BlockedByQueuedChildren, queuedIds.Count, 0);
            }

            await _context.Tasks
                .Where(t => queuedIds.Contains(t.Id))
                .ExecuteUpdateAsync(s => s
                    .SetProperty(t => t.Status, TaskStatus.Idle)
                    .SetProperty(t => t.BlockedByTaskId, (string?)null), ct);
        }

        // Terminal children (Done/Failed/Cancelled) stay attached to the parent even
        // though its PlanningPhase will be reset to None. The lineage is preserved as
        // historical context; the UI nests them under their parent regardless of phase.

        // Idle children created during this planning session are dropped.
        await _context.Tasks
            .Where(t => t.ParentTaskId == parentId
                        && t.Status == TaskStatus.Idle
                        && t.PlanningPhase == PlanningPhase.None)
            .ExecuteDeleteAsync(ct);

        await _context.Tasks
            .Where(t => t.Id == parentId)
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Idle)
                .SetProperty(t => t.PlanningPhase, PlanningPhase.None)
                .SetProperty(t => t.PlanningSessionId, (string?)null)
                .SetProperty(t => t.PlanningSessionToken, (string?)null)
                .SetProperty(t => t.PlanningFinalizedAt, (DateTime?)null), ct);

        await tx.CommitAsync(ct);
        return new DiscardPlanningOutcome(DiscardPlanningResult.Discarded, queuedIds.Count, 0);
    }

    /// <summary>
    /// Dequeues child tasks whose parent is missing or no longer in a planning phase:
    /// sets <c>Status</c> from <c>Queued</c> to <c>Idle</c> and clears
    /// <c>BlockedByTaskId</c>. <c>ParentTaskId</c> stays intact — the child remains
    /// part of its (former) planning chain for historical context. Returns the
    /// number of rows dequeued. Idempotent.
    /// </summary>
    internal async Task<int> DequeueOrphanedChildrenAsync(CancellationToken ct = default)
    {
        var orphanIds = await _context.Tasks
            .Where(t => t.ParentTaskId != null && t.Status == TaskStatus.Queued)
            // Agent-suggested improvement children (CreatedBy == ParentTaskId) legitimately
            // queue under a non-planning parent — they are not orphaned planning-chain members.
            .Where(t => t.CreatedBy == null || t.CreatedBy != t.ParentTaskId)
            .Where(t => !_context.Tasks.Any(p =>
                p.Id == t.ParentTaskId && p.PlanningPhase != PlanningPhase.None))
            .Select(t => t.Id)
            .ToListAsync(ct);

        if (orphanIds.Count == 0) return 0;

        return await _context.Tasks
            .Where(t => orphanIds.Contains(t.Id))
            .ExecuteUpdateAsync(s => s
                .SetProperty(t => t.Status, TaskStatus.Idle)
                .SetProperty(t => t.BlockedByTaskId, (string?)null), ct);
    }

    /// <summary>
    /// Restores a planning-session lineage that lost its <c>parent_task_id</c> links.
    /// Given a candidate parent task and a single unambiguous orphan chain in the
    /// same list (linked via <c>BlockedByTaskId</c>), re-attaches the chain members
    /// to the parent, marks the parent as <c>Finalized</c>, and dequeues queued
    /// chain members. No-op if conditions are not met. Returns the number of
    /// re-attached children (0 if skipped).
    /// </summary>
    internal async Task<int> RestorePlanningLineageAsync(string parentId, CancellationToken ct = default)
    {
        var parent = await _context.Tasks.AsNoTracking()
            .FirstOrDefaultAsync(t => t.Id == parentId, ct);
        if (parent is null) return 0;
        if (parent.PlanningPhase != PlanningPhase.None) return 0;
        if (parent.Status is TaskStatus.Done or TaskStatus.Failed or TaskStatus.Cancelled) return 0;

        // Candidates: unattached tasks in the same list, excluding the parent itself.
        var candidates = await _context.Tasks.AsNoTracking()
            .Where(t => t.ListId == parent.ListId && t.ParentTaskId == null && t.Id != parent.Id)
            .ToListAsync(ct);

        // A chain is a maximal linear sequence linked via BlockedByTaskId. Find heads
        // (BlockedByTaskId == null) that have at least one successor.
        var bySource = candidates
            .Where(c => c.BlockedByTaskId != null)
            .ToLookup(c => c.BlockedByTaskId!);

        var heads = candidates
            .Where(c => c.BlockedByTaskId == null && bySource[c.Id].Any())
            .ToList();

        // Bail unless exactly one chain anchors a successor — anything else is
        // ambiguous and we refuse to guess.
        if (heads.Count != 1) return 0;

        var chain = new List<TaskEntity> { heads[0] };
        var current = heads[0];
        while (true)
        {
            var next = bySource[current.Id].FirstOrDefault();
            if (next is null) break;
            chain.Add(next);
            current = next;
        }

        var chainIds = chain.Select(c => c.Id).ToList();

        await _context.Tasks
            .Where(t => t.Id == parentId)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.PlanningPhase, PlanningPhase.Finalized), ct);

        await _context.Tasks
            .Where(t => chainIds.Contains(t.Id))
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.ParentTaskId, (string?)parentId), ct);

        // Dequeue queued chain members; blocked_by stays intact so chain order is
        // preserved for manual re-queueing.
        await _context.Tasks
            .Where(t => chainIds.Contains(t.Id) && t.Status == TaskStatus.Queued)
            .ExecuteUpdateAsync(s => s.SetProperty(t => t.Status, TaskStatus.Idle), ct);

        return chainIds.Count;
    }

    #endregion
}