src/Infrastructure/TelegramSaaS.API/Controllers/GpuRegistrationController.cs

using System;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Mvc;
using MongoDB.Driver;
using TelegramSaaS.Domain.Entities;
using TelegramSaaS.Infrastructure.Persistence;

namespace TelegramSaaS.API.Controllers;

[ApiController]
[Route("api/v1/gpu")]
public class GpuRegistrationController : ControllerBase
{
    private readonly MongoDbContext _dbContext;
    private const string CollectionName = "GpuServers";
    private static readonly TimeSpan HeartbeatTimeout = TimeSpan.FromHours(6);

    public GpuRegistrationController(MongoDbContext dbContext)
    {
        _dbContext = dbContext;
    }

    [HttpPost("register")]
    public async Task<IActionResult> Register([FromBody] GpuRegisterRequest request)
    {
        if (string.IsNullOrWhiteSpace(request.Url) || string.IsNullOrWhiteSpace(request.Label))
            return BadRequest(new { error = "Both 'url' and 'label' are required." });

        var label = request.Label.Trim().ToLowerInvariant();
        var collection = _dbContext.GetCollection<GpuServer>(CollectionName);

        var filter = Builders<GpuServer>.Filter.Eq(s => s.Label, label);
        var existing = await collection.Find(filter).FirstOrDefaultAsync();

        if (existing != null)
        {
            var update = Builders<GpuServer>.Update
                .Set(s => s.Url, request.Url.Trim())
                .Set(s => s.LastHeartbeat, DateTime.UtcNow)
                .Set(s => s.IsActive, true)
                .Set(s => s.UpdatedAt, DateTime.UtcNow);

            await collection.UpdateOneAsync(filter, update);
        }
        else
        {
            var server = new GpuServer
            {
                Url = request.Url.Trim(),
                Label = label,
                LastHeartbeat = DateTime.UtcNow,
                IsActive = true
            };
            await collection.InsertOneAsync(server);
        }

        return Ok(new { message = $"GPU server '{label}' registered successfully.", url = request.Url.Trim() });
    }

    [HttpPost("heartbeat")]
    public async Task<IActionResult> Heartbeat([FromBody] GpuHeartbeatRequest request)
    {
        if (string.IsNullOrWhiteSpace(request.Url))
            return BadRequest(new { error = "'url' is required." });

        var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
        var filter = Builders<GpuServer>.Filter.Eq(s => s.Url, request.Url.Trim());

        var update = Builders<GpuServer>.Update
            .Set(s => s.LastHeartbeat, DateTime.UtcNow)
            .Set(s => s.IsActive, true)
            .Set(s => s.UpdatedAt, DateTime.UtcNow);

        var result = await collection.UpdateOneAsync(filter, update);

        if (result.MatchedCount == 0)
            return NotFound(new { error = "GPU server not found." });

        return Ok(new { message = "Heartbeat updated." });
    }

    [HttpGet("urls")]
    public async Task<IActionResult> GetActiveUrls()
    {
        var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
        var cutoff = DateTime.UtcNow.Add(-HeartbeatTimeout);

        var filter = Builders<GpuServer>.Filter.And(
            Builders<GpuServer>.Filter.Eq(s => s.IsActive, true),
            Builders<GpuServer>.Filter.Gt(s => s.LastHeartbeat, cutoff)
        );

        var servers = await collection.Find(filter).ToListAsync();

        var result = servers.Select(s => new
        {
            url = s.Url,
            label = s.Label,
            lastHeartbeat = s.LastHeartbeat
        });

        return Ok(new { servers = result });
    }

    [HttpDelete("unregister")]
    public async Task<IActionResult> Unregister([FromBody] GpuUnregisterRequest request)
    {
        if (string.IsNullOrWhiteSpace(request.Url))
            return BadRequest(new { error = "'url' is required." });

        var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
        var filter = Builders<GpuServer>.Filter.Eq(s => s.Url, request.Url.Trim());

        var result = await collection.DeleteOneAsync(filter);

        if (result.DeletedCount == 0)
            return NotFound(new { error = "GPU server not found." });

        return Ok(new { message = "GPU server unregistered." });
    }
}

// Request DTOs
public record GpuRegisterRequest(string Url, string Label);
public record GpuHeartbeatRequest(string Url);
public record GpuUnregisterRequest(string Url);