ClawdDotNet/src/ClawdDotNet.Tools.TelegramClient/TelegramClientManager.cs

using ClawdDotNet.Core.Config;
using Microsoft.Extensions.Logging;
using TL;
using WTelegram;

namespace ClawdDotNet.Tools.TelegramClient;

public sealed class TelegramClientManager : IAsyncDisposable
{
    private Client? _client;
    private User? _self;
    private readonly SemaphoreSlim _gate = new(1, 1);
    private readonly ILogger _logger;

    private readonly int _apiId;
    private readonly string _apiHash;
    private readonly string _phoneNumber;
    private readonly string _sessionPath;
    private readonly string? _2faPassword;

    private readonly Dictionary<long, User> _users = new();
    private readonly Dictionary<long, ChatBase> _chats = new();

    public event Func<string, Task<string>>? OnLoginCodeRequired;
    public event Func<Task<string>>? On2FAPasswordRequired;

    public bool IsConnected => _client?.User != null;
    public User? Self => _self;

    public TelegramClientManager(
        InstanceConfig config,
        string instancePath,
        ILogger logger)
    {
        _logger = logger;

        var tgConfig = config.TelegramClient
            ?? throw new InvalidOperationException("TelegramClient config missing in InstanceConfig");

        _apiId = tgConfig.ApiId;
        _apiHash = tgConfig.ApiHash;
        _phoneNumber = tgConfig.PhoneNumber;
        _sessionPath = Path.Combine(instancePath, "sessions", $"telegram_{config.InstanceId}.session");
        _2faPassword = tgConfig.Password2FA;
    }

    public async Task ConnectAsync(CancellationToken ct)
    {
        var sessionDir = Path.GetDirectoryName(_sessionPath)!;
        Directory.CreateDirectory(sessionDir);

        // #794: A login can be started again after it failed; the old client goes first.
        _client?.Dispose();
        _client = new Client(ConfigCallback);

        Helpers.Log = (lvl, msg) =>
            _logger.Log((LogLevel)lvl, "WTelegram: {Message}", msg);

        _self = await _client.LoginUserIfNeeded();
        _logger.LogInformation(
            "Telegram: logged in as {Name} (id {Id})",
            _self.first_name, _self.id);
    }

    private string? ConfigCallback(string what) => what switch
    {
        "api_id" => _apiId.ToString(),
        "api_hash" => _apiHash,
        "phone_number" => _phoneNumber,
        "session_pathname" => _sessionPath,

        "verification_code" => OnLoginCodeRequired != null
            ? OnLoginCodeRequired("Please enter Telegram verification code:").Result
            : throw new InvalidOperationException(
                "Verification code required but no UI handler registered. " +
                "Connect OnLoginCodeRequired to prompt the user."),

        "password" => _2faPassword
            ?? (On2FAPasswordRequired != null
                ? On2FAPasswordRequired().Result
                : throw new InvalidOperationException(
                    "2FA password required but not configured.")),

        _ => null
    };

    public async Task<Messages_Dialogs> GetAllDialogsAsync(CancellationToken ct)
    {
        await _gate.WaitAsync(ct);
        try
        {
            var dialogs = await _client!.Messages_GetAllDialogs();
            dialogs.CollectUsersChats(_users, _chats);
            return dialogs;
        }
        finally { _gate.Release(); }
    }

    public async Task<Messages_Chats> GetAllChatsAsync(CancellationToken ct)
    {
        await _gate.WaitAsync(ct);
        try { return await _client!.Messages_GetAllChats(); }
        finally { _gate.Release(); }
    }

    public async Task<Messages_MessagesBase> GetMessagesAsync(
        InputPeer peer, int minId = 0, int limit = 50, CancellationToken ct = default)
    {
        await _gate.WaitAsync(ct);
        try
        {
            return await _client!.Messages_GetHistory(
                peer, offset_id: 0, offset_date: default,
                add_offset: 0, limit: limit, max_id: 0, min_id: minId, hash: 0);
        }
        finally { _gate.Release(); }
    }

    public async Task<Contacts_ResolvedPeer> ResolveUsernameAsync(string username, CancellationToken ct)
    {
        await _gate.WaitAsync(ct);
        try
        {
            var resolved = await _client!.Contacts_ResolveUsername(username.TrimStart('@'));
            foreach (var u in resolved.users.Values)
                _users[u.id] = u;
            foreach (var c in resolved.chats.Values)
                _chats[c.ID] = c;
            return resolved;
        }
        finally { _gate.Release(); }
    }

    public InputPeer? GetInputPeerFromCache(long chatId)
    {
        if (_users.TryGetValue(chatId, out var user))
            return user;
        if (_chats.TryGetValue(chatId, out var chat))
            return chat;
        return null;
    }

    public async ValueTask DisposeAsync()
    {
        _client?.Dispose();
        _gate.Dispose();
    }
}