PlotDirector/PlotLine/Hubs/WordCompanionFollowHub.cs

using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.SignalR;
using PlotLine.Data;
using PlotLine.Models;
using PlotLine.Services;
using System.Security.Claims;

namespace PlotLine.Hubs;

[Authorize]
public sealed class WordCompanionFollowHub(
    IWordCompanionPresenceService presence,
    IWritingWorkspacePresenceService writingPresence,
    IOnboardingRepository onboarding,
    IManuscriptScanPreviewStore scanStore,
    IOnboardingService onboardingService,
    IOnboardingBuildRepository onboardingBuilds,
    ILogger<WordCompanionFollowHub> logger) : Hub
{
    public override async Task OnConnectedAsync()
    {
        if (TryGetUserId(out var userId))
        {
            await Groups.AddToGroupAsync(Context.ConnectionId, PresenceGroup(userId));
        }

        await base.OnConnectedAsync();
    }

    public async Task<WordCompanionPresenceStatus> WatchCompanionPresence()
    {
        var userId = RequireUserId();
        await Groups.AddToGroupAsync(Context.ConnectionId, PresenceGroup(userId));
        return await presence.GetStatusAsync(userId);
    }

    public async Task<WordCompanionPresenceStatus> RegisterCompanion(WordCompanionPresenceRegistration registration)
    {
        var userId = RequireUserId();
        if (registration.UserID.HasValue && registration.UserID.Value != userId)
        {
            throw new HubException("The Companion is signed in as a different user.");
        }

        var status = await presence.RegisterAsync(userId, Context.ConnectionId, registration);
        logger.LogInformation(
            "Word Companion registered. UserID={UserID} ProjectID={ProjectID} BookID={BookID} DocumentOpen={DocumentOpen}",
            userId,
            registration.LinkedProjectID,
            registration.LinkedBookID,
            registration.DocumentOpen);
        await Clients.Group(PresenceGroup(userId)).SendAsync("WordCompanionPresenceChanged", status);
        await BroadcastWritingWorkspaceStatusAsync(userId, registration.LinkedProjectID, registration.LinkedBookID);
        return status;
    }

    public async Task<WordCompanionPresenceStatus> CompanionHeartbeat(WordCompanionPresenceHeartbeat heartbeat)
    {
        var userId = RequireUserId();
        var status = await presence.HeartbeatAsync(userId, Context.ConnectionId, heartbeat);
        logger.LogDebug(
            "Word Companion heartbeat. UserID={UserID} ProjectID={ProjectID} BookID={BookID} DocumentOpen={DocumentOpen}",
            userId,
            heartbeat.LinkedProjectID,
            heartbeat.LinkedBookID,
            heartbeat.DocumentOpen);
        await Clients.Group(PresenceGroup(userId)).SendAsync("WordCompanionPresenceChanged", status);
        await BroadcastWritingWorkspaceStatusAsync(userId, heartbeat.LinkedProjectID, heartbeat.LinkedBookID);
        return status;
    }

    public async Task<WritingWorkspaceStatus> WatchWritingWorkspace(WritingWorkspaceScope scope)
    {
        var userId = RequireUserId();
        await Groups.AddToGroupAsync(Context.ConnectionId, PresenceGroup(userId));
        logger.LogInformation(
            "Writing workspace watched. UserID={UserID} ProjectID={ProjectID} BookID={BookID}",
            userId,
            scope.ProjectID,
            scope.BookID);
        return await WritingWorkspaceStatusAsync(userId, scope);
    }

    public async Task<WritingWorkspaceStatus> RegisterWritingWindow(WritingWindowRegistration registration)
    {
        var userId = RequireUserId();
        await Groups.AddToGroupAsync(Context.ConnectionId, PresenceGroup(userId));
        var status = await writingPresence.RegisterAsync(userId, Context.ConnectionId, registration, await presence.GetStatusAsync(userId));
        logger.LogInformation(
            "Writing window registered. UserID={UserID} Role={Role} ProjectID={ProjectID} BookID={BookID} Route={Route}",
            userId,
            registration.WritingRole,
            registration.ProjectID,
            registration.BookID,
            registration.CurrentRoute);
        await Clients.Group(PresenceGroup(userId)).SendAsync("WritingWorkspacePresenceChanged", status);
        return status;
    }

    public async Task<WritingWorkspaceStatus> WritingWindowHeartbeat(WritingWindowHeartbeat heartbeat)
    {
        var userId = RequireUserId();
        var status = await writingPresence.HeartbeatAsync(userId, Context.ConnectionId, heartbeat, await presence.GetStatusAsync(userId));
        await Clients.Group(PresenceGroup(userId)).SendAsync("WritingWorkspacePresenceChanged", status);
        return status;
    }

    public async Task<WritingWorkspaceStatus> RequestWritingWindowFocus(WritingWorkspaceScope scope, string role)
    {
        var userId = RequireUserId();
        var connectionId = await writingPresence.GetConnectionIdForRoleAsync(userId, scope, role);
        if (!string.IsNullOrWhiteSpace(connectionId))
        {
            await Clients.Client(connectionId).SendAsync("WritingWorkspaceFocusRequested", new
            {
                WritingRole = WritingWorkspaceRoles.Normalise(role),
                scope.ProjectID,
                scope.BookID
            });
        }

        return await WritingWorkspaceStatusAsync(userId, scope);
    }

    public async Task<ManuscriptScanState> StartOnboardingManuscriptScan()
    {
        var userId = RequireUserId();
        var state = await onboarding.GetAsync(userId);
        if (state is null
            || !IsMicrosoftWordState(state)
            || !state.ProjectID.HasValue
            || !state.BookID.HasValue)
        {
            throw new HubException("Finish choosing a Project and Book before scanning.");
        }

        var presenceStatus = await presence.GetStatusAsync(userId);
        if (!presenceStatus.IsConnected)
        {
            throw new HubException("The Word Companion is not connected. Open Word and try again.");
        }

        if (!presenceStatus.DocumentOpen)
        {
            throw new HubException("Open your manuscript in Word before scanning.");
        }

        var companionConnectionId = await presence.GetCompanionConnectionIdAsync(userId);
        if (string.IsNullOrWhiteSpace(companionConnectionId))
        {
            throw new HubException("The Word Companion connection is not ready. Try reconnecting the Companion.");
        }

        var scanState = await scanStore.StartAsync(userId, state.UserOnboardingStateID, state.ProjectID.Value, state.BookID.Value);
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingScanStateChanged", scanState);
        await Clients.Client(companionConnectionId).SendAsync("ScanCurrentDocument", new ManuscriptScanCommand
        {
            UserID = userId,
            OnboardingID = state.UserOnboardingStateID,
            ProjectID = state.ProjectID.Value,
            BookID = state.BookID.Value,
            RequestedUtc = DateTime.UtcNow
        });
        return scanState;
    }

    public async Task<ManuscriptScanState> ReportOnboardingScanProgress(ManuscriptScanProgress progress)
    {
        var userId = RequireUserId();
        var scanState = await scanStore.ProgressAsync(userId, progress);
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingScanProgress", scanState);
        return scanState;
    }

    public async Task<ManuscriptScanState> CompleteOnboardingScan(ManuscriptScanPreview preview)
    {
        var userId = RequireUserId();
        var state = await onboarding.GetAsync(userId);
        if (state is not null)
        {
            preview = NormalisePreviewContext(preview, userId, state);
        }

        if (state is null
            || state.UserOnboardingStateID != preview.OnboardingID
            || state.ProjectID != preview.ProjectID
            || state.BookID != preview.BookID)
        {
            throw new HubException("The scan result does not match your onboarding setup.");
        }

        var scanState = await scanStore.CompleteAsync(userId, preview);
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingScanCompleted", scanState);
        return scanState;
    }

    public async Task<ManuscriptScanState> FailOnboardingScan(ManuscriptScanFailure failure)
    {
        var userId = RequireUserId();
        var scanState = await scanStore.FailAsync(userId, failure.OnboardingID, failure.Message);
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingScanFailed", scanState);
        return scanState;
    }

    public async Task<OnboardingManuscriptBuildResult> StartOnboardingProjectBuild(Guid previewId)
    {
        var userId = RequireUserId();
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildProgress", new OnboardingBuildProgress
        {
            Status = "Running",
            Message = "Creating chapters...",
            PercentComplete = 0,
            PreviewID = previewId
        });

        OnboardingManuscriptBuildResult? result;
        try
        {
            result = await onboardingService.BuildApprovedStructureAsync(
                previewId,
                progress => Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildProgress", progress));
        }
        catch (InvalidOperationException ex)
        {
            await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildProgress", new OnboardingBuildProgress
            {
                Status = "Failed",
                Message = ex.Message,
                PercentComplete = 25,
                PreviewID = previewId
            });
            throw new HubException(ex.Message);
        }

        if (result is null)
        {
            throw new HubException("The reviewed scan could not be found.");
        }

        var companionConnectionId = await presence.GetCompanionConnectionIdAsync(userId);
        if (!string.IsNullOrWhiteSpace(companionConnectionId))
        {
            await onboardingBuilds.MarkMarkerSyncSentAsync(result.BuildID);
            await Clients.Client(companionConnectionId).SendAsync("UpdateOnboardingBuildMarkers", result);
        }
        else
        {
            await onboardingBuilds.MarkMarkerSyncFailedAsync(result.BuildID, "The Word Companion disconnected before marker sync could be sent.");
            await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildProgress", new OnboardingBuildProgress
            {
                Status = "Failed",
                Message = "The Word Companion disconnected before PlotDirector could update and save Word markers.",
                PercentComplete = 90,
                ChaptersCreated = result.ChaptersCreated,
                ScenesCreated = result.ScenesCreated,
                CharactersCreated = result.CharactersCreated + result.CharactersReused,
                SceneAppearancesCreated = result.SceneAppearancesCreated,
                PreviewID = previewId
            });
        }

        return result;
    }

    public async Task CompleteOnboardingBuildMarkerSync(int buildId)
    {
        var userId = RequireUserId();
        await onboardingBuilds.MarkMarkerSyncCompletedAsync(buildId);
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildProgress", new OnboardingBuildProgress
        {
            Status = "Complete",
            Message = "Manuscript structure has been created and saved. Word may now be closed; keep this browser workflow open to review characters.",
            PercentComplete = 100
        });
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildMarkerSyncCompleted", new { buildId });
    }

    public async Task FailOnboardingBuildMarkerSync(int buildId, string message)
    {
        var userId = RequireUserId();
        await onboardingBuilds.MarkMarkerSyncFailedAsync(buildId, message);
        await Clients.Group(PresenceGroup(userId)).SendAsync("OnboardingBuildProgress", new OnboardingBuildProgress
        {
            Status = "Failed",
            Message = string.IsNullOrWhiteSpace(message)
                ? "Word marker sync failed. Keep Word open and try again."
                : message,
            PercentComplete = 90
        });
    }

    public override async Task OnDisconnectedAsync(Exception? exception)
    {
        int? disconnectedUserId = null;
        var status = await presence.DisconnectAsync(Context.ConnectionId);
        if (status is not null)
        {
            disconnectedUserId = status.UserID;
            await Clients.Group(PresenceGroup(status.UserID)).SendAsync("WordCompanionPresenceChanged", status);
            await BroadcastWritingWorkspaceStatusAsync(status.UserID, status.LinkedProjectID, status.LinkedBookID);
            var failedScans = await scanStore.FailRunningForUserAsync(
                status.UserID,
                "The Word Companion disconnected before the scan finished. Reopen Word and try again.");
            foreach (var failedScan in failedScans)
            {
                await Clients.Group(PresenceGroup(status.UserID)).SendAsync("OnboardingScanFailed", failedScan);
            }
        }

        var writingStatus = await writingPresence.DisconnectAsync(
            Context.ConnectionId,
            disconnectedUserId.HasValue ? await presence.GetStatusAsync(disconnectedUserId.Value) : null);
        if (writingStatus is not null)
        {
            await Clients.Group(PresenceGroup(writingStatus.UserID)).SendAsync("WritingWorkspacePresenceChanged", writingStatus);
        }

        await base.OnDisconnectedAsync(exception);
    }

    private async Task<WritingWorkspaceStatus> WritingWorkspaceStatusAsync(int userId, WritingWorkspaceScope scope)
        => await writingPresence.GetStatusAsync(userId, scope, await presence.GetStatusAsync(userId));

    private async Task BroadcastWritingWorkspaceStatusAsync(int userId, int? projectId, int? bookId)
    {
        var workspaceStatus = await WritingWorkspaceStatusAsync(userId, new WritingWorkspaceScope
        {
            ProjectID = projectId,
            BookID = bookId
        });
        await Clients.Group(PresenceGroup(userId)).SendAsync("WritingWorkspacePresenceChanged", workspaceStatus);
    }

    private int RequireUserId()
        => TryGetUserId(out var userId) ? userId : throw new HubException("Sign in to use the Word Companion.");

    private bool TryGetUserId(out int userId)
        => int.TryParse(Context.User?.FindFirstValue(ClaimTypes.NameIdentifier), out userId);

    private static string PresenceGroup(int userId) => $"word-companion-presence:{userId}";

    private static bool IsMicrosoftWordState(UserOnboardingState state)
        => string.Equals(state.WritingSoftware, WritingSoftwareValues.MicrosoftWord, StringComparison.Ordinal)
            || state.UsesWordCompanion == true;

    private static ManuscriptScanPreview NormalisePreviewContext(ManuscriptScanPreview preview, int userId, UserOnboardingState state)
    {
        return new()
        {
            PreviewID = preview.PreviewID,
            UserID = preview.UserID == 0 ? userId : preview.UserID,
            OnboardingID = preview.OnboardingID == 0 ? state.UserOnboardingStateID : preview.OnboardingID,
            ProjectID = preview.ProjectID == 0 && state.ProjectID.HasValue ? state.ProjectID.Value : preview.ProjectID,
            BookID = preview.BookID == 0 && state.BookID.HasValue ? state.BookID.Value : preview.BookID,
            Source = preview.Source,
            DocumentTitle = preview.DocumentTitle,
            CompanionDocumentIdentifier = preview.CompanionDocumentIdentifier,
            TotalWordCount = preview.TotalWordCount,
            ChapterCount = preview.ChapterCount,
            SceneCount = 0,
            CharacterCandidateCount = 0,
            CreatedUtc = preview.CreatedUtc,
            Chapters = preview.Chapters,
            Scenes = [],
            CharacterCandidates = []
        };
    }
}