TermRemoteCtl/apps/windows_agent/src/TermRemoteCtl.Agent/Terminal/PowerShellSessionHost.cs

using System.Collections.Concurrent;
using System.Runtime.Versioning;
using TermRemoteCtl.Agent.Sessions;

namespace TermRemoteCtl.Agent.Terminal;

[SupportedOSPlatform("windows")]
internal sealed class PowerShellSessionHost : ISessionHost, IAsyncDisposable
{
    private readonly IConPtySessionFactory _sessionFactory;
    private readonly SessionRegistry _sessionRegistry;
    private readonly ConcurrentDictionary<string, IConPtySession> _sessions = new(StringComparer.Ordinal);

    public PowerShellSessionHost(IConPtySessionFactory sessionFactory, SessionRegistry sessionRegistry)
    {
        _sessionFactory = sessionFactory;
        _sessionRegistry = sessionRegistry;
    }

    public event EventHandler<TerminalOutputEventArgs>? OutputReceived;

    public async Task StartAsync(string sessionId, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(sessionId);
        ConPtyInterop.EnsureSupported();

        if (_sessions.TryGetValue(sessionId, out var existingSession))
        {
            if (await existingSession.IsAliveAsync(cancellationToken).ConfigureAwait(false))
            {
                return;
            }

            if (_sessions.TryRemove(new KeyValuePair<string, IConPtySession>(sessionId, existingSession)))
            {
                existingSession.OutputReceived -= HandleSessionOutput;
                await existingSession.DisposeAsync().ConfigureAwait(false);
            }
        }

        if (!_sessionRegistry.TryGet(sessionId, out var record) || record is null)
        {
            throw new KeyNotFoundException($"Session '{sessionId}' is not registered.");
        }

        var session = _sessionFactory.Create(sessionId, record.WorkingDirectory);
        if (!_sessions.TryAdd(sessionId, session))
        {
            await session.DisposeAsync().ConfigureAwait(false);
            return;
        }

        session.OutputReceived += HandleSessionOutput;

        try
        {
            await session.StartAsync(cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            session.OutputReceived -= HandleSessionOutput;
            _sessions.TryRemove(sessionId, out _);
            await session.DisposeAsync().ConfigureAwait(false);
            throw;
        }
    }

    public async Task WriteInputAsync(string sessionId, string input, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(sessionId);
        ArgumentNullException.ThrowIfNull(input);

        if (!_sessions.TryGetValue(sessionId, out var session))
        {
            throw new KeyNotFoundException($"Session '{sessionId}' is not running.");
        }

        await session.WriteInputAsync(input, cancellationToken).ConfigureAwait(false);
    }

    public async Task ResizeAsync(string sessionId, int columns, int rows, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(sessionId);

        if (!_sessions.TryGetValue(sessionId, out var session))
        {
            throw new KeyNotFoundException($"Session '{sessionId}' is not running.");
        }

        await session.ResizeAsync(columns, rows, cancellationToken).ConfigureAwait(false);
    }

    public async Task StopAsync(string sessionId, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(sessionId);
        cancellationToken.ThrowIfCancellationRequested();

        if (!_sessions.TryRemove(sessionId, out var session))
        {
            return;
        }

        session.OutputReceived -= HandleSessionOutput;
        await session.DisposeAsync().ConfigureAwait(false);
    }

    public async ValueTask DisposeAsync()
    {
        foreach (var session in _sessions.Values)
        {
            session.OutputReceived -= HandleSessionOutput;
            await session.DisposeAsync().ConfigureAwait(false);
        }

        _sessions.Clear();
    }

    private void HandleSessionOutput(object? sender, TerminalOutputEventArgs args)
    {
        _ = _sessionRegistry.AppendOutputAsync(args.SessionId, args.Chunk, CancellationToken.None);
        OutputReceived?.Invoke(this, args);
    }
}