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 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(CollectionName); var filter = Builders.Filter.Eq(s => s.Label, label); var existing = await collection.Find(filter).FirstOrDefaultAsync(); if (existing != null) { var update = Builders.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 Heartbeat([FromBody] GpuHeartbeatRequest request) { if (string.IsNullOrWhiteSpace(request.Url)) return BadRequest(new { error = "'url' is required." }); var collection = _dbContext.GetCollection(CollectionName); var filter = Builders.Filter.Eq(s => s.Url, request.Url.Trim()); var update = Builders.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 GetActiveUrls() { var collection = _dbContext.GetCollection(CollectionName); var cutoff = DateTime.UtcNow.Add(-HeartbeatTimeout); var filter = Builders.Filter.And( Builders.Filter.Eq(s => s.IsActive, true), Builders.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 Unregister([FromBody] GpuUnregisterRequest request) { if (string.IsNullOrWhiteSpace(request.Url)) return BadRequest(new { error = "'url' is required." }); var collection = _dbContext.GetCollection(CollectionName); var filter = Builders.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);