Program.cs

using System.Diagnostics;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using GithubActionsOrchestrator.Auth;
using GithubActionsOrchestrator.CloudControllers;
using GithubActionsOrchestrator.Database;
using GithubActionsOrchestrator.GitHub;
using GithubActionsOrchestrator.Models;
using Microsoft.AspNetCore.Mvc;
using Microsoft.EntityFrameworkCore;
using OpenTelemetry.Exporter;
using OpenTelemetry.Resources;
using OpenTelemetry.Trace;
using Prometheus;
using Serilog;
using Serilog.Events;

namespace GithubActionsOrchestrator;

public class Program
{
    public static AutoScalerConfiguration Config = new();

    internal const string ServiceName = "github-actions-orchestrator";
    internal static readonly ActivitySource OrchestratorActivitySource = new(ServiceName);

    private static readonly Counter ProcessedJobCount = Metrics
        .CreateCounter("github_autoscaler_jobs_processed", "Number of processed jobs", labelNames: ["org", "size"]);

    private static readonly Counter QueuedJobCount = Metrics
        .CreateCounter("github_autoscaler_jobs_queued", "Number of queued jobs", labelNames: ["org", "size"]);

    private static readonly Counter PickedJobCount = Metrics
        .CreateCounter("github_autoscaler_jobs_picked", "Number of jobs picked up by a runner",
            labelNames: ["org", "size"]);

    private static readonly Counter MachineFailedCount = Metrics
        .CreateCounter("github_autoscaler_machine_failed", "Number of machines failed to provision",
            labelNames: ["org", "size"]);
    
    private static readonly Counter MachineSuccessCount = Metrics
        .CreateCounter("github_autoscaler_machine_success", "Number of machines provisioned fine",
            labelNames: ["org", "size"]);
    
    private static readonly Counter TotalMachineTime = Metrics
        .CreateCounter("github_autoscaler_total_machine_time", "Number of seconds machines were alive",
            labelNames: ["org", "size"]);

    private static readonly Counter CancelledJobCount = Metrics
        .CreateCounter("github_autoscaler_jobs_cancelled", "Number of jobs cancelled before runner assignment",
            labelNames: ["org", "repo", "size", "profile", "arch"]);

    public static void Main(string[] args)
    {
        //Init GlitchTip
        SentrySdk.Init(options =>
        {
            options.Dsn = "https://32e67fe320454af79c685a59a3502116@glitchtip.ethquokkaops.io/1";
        });
        
        // Set up logging
        Log.Logger = new LoggerConfiguration()
            .MinimumLevel.Debug()
            .MinimumLevel.Override("Microsoft", LogEventLevel.Warning)
            .MinimumLevel.Override("System", LogEventLevel.Warning)
            .Enrich.FromLogContext()
            .WriteTo.Console()
            .CreateLogger();

        if (!LoadConfiguration())
        {
            Log.Error("Unable to do initial config load. Aborting.");
            return;
        }

        if (Config.RunnerPrefix == "ghr" &&
            Environment.GetEnvironmentVariable("ALLOW_PRODUCTION_PREFIX") != "true")
        {
            Log.Error("SAFETY: RunnerPrefix is 'ghr' (production). " +
                      "Set ALLOW_PRODUCTION_PREFIX=true to confirm, or use 'ghr-dev' for local dev.");
            return;
        }

        // Prepare metrics
        using KestrelMetricServer server = new(port: 9000);
        server.Start();
        Log.Information("Metrics server listening on port 9000");

        // Init count metrics
        foreach (GithubTargetConfiguration org in Config.TargetConfigs)
        {
            foreach (MachineSize ms in Config.Sizes)
            {
                TotalMachineTime.Labels(org.Name, ms.Name).IncTo(0);
                PickedJobCount.Labels(org.Name, ms.Name).IncTo(0);
                QueuedJobCount.Labels(org.Name, ms.Name).IncTo(0);
                ProcessedJobCount.Labels(org.Name, ms.Name).IncTo(0);
            }
        }

        // Database migration
        Log.Information("Running Database migrations...");
        using (var db = new ActionsRunnerContext())
        {
            db.Database.Migrate();
        }
        Log.Information("Migrations complete.");


        WebApplicationBuilder builder = WebApplication.CreateBuilder(args);
        builder.Services.AddSerilog();
        builder.Services.AddOpenTelemetry()
            .ConfigureResource(resource => resource.AddService(serviceName: ServiceName))
            .WithTracing(tracing =>
            {
                tracing
                    .AddSource(ServiceName)
                    .AddSource("Npgsql")
                    .AddAspNetCoreInstrumentation()
                    .AddHttpClientInstrumentation()
                    .AddOtlpExporter(opt =>
                    {
                        opt.Protocol = OtlpExportProtocol.HttpProtobuf;
                        var otlpUser = Environment.GetEnvironmentVariable("OTEL_EXPORTER_OTLP_AUTH_USER");
                        var otlpPass = Environment.GetEnvironmentVariable("OTEL_EXPORTER_OTLP_AUTH_PASSWORD");
                        if (!string.IsNullOrEmpty(otlpUser) && !string.IsNullOrEmpty(otlpPass))
                        {
                            var credentials = Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes($"{otlpUser}:{otlpPass}"));
                            opt.Headers = $"Authorization=Basic {credentials}";
                        }
                    })
                    .AddProcessor(new Pyroscope.OpenTelemetry.PyroscopeSpanProcessor());
            });
        builder.Services.AddSingleton<RunnerQueue>();
        builder.Services.AddHostedService<PoolManager>();

        // Add services to the container.
        

        // Learn more about configuring Swagger/OpenAPI at https://aka.ms/aspnetcore/swashbuckle
        builder.Services.AddCors(options =>
        {
            options.AddPolicy("AllowAll", builder =>
            {
                builder
                    .AllowAnyOrigin()    // Allow all origins
                    .AllowAnyMethod()    // Allow all HTTP methods
                    .AllowAnyHeader();   // Allow all headers
            });
        });
        builder.Services.AddControllers().AddJsonOptions(options =>
         {
             options.JsonSerializerOptions.ReferenceHandler = ReferenceHandler.IgnoreCycles;
         });

        // Identity + authorization for mutating dashboard endpoints. The provider is chosen
        // by config: Warpgate wins if both are enabled. When the active provider is
        // disabled, mutating endpoints are gated by the API key only.
        if (Config.WarpgateAuth.Enabled)
        {
            if (Config.TeleportAuth.Enabled)
                Log.Warning("Both WarpgateAuth and TeleportAuth are enabled - Warpgate takes precedence.");
            builder.AddWarpgateAuth(Config.WarpgateAuth);
            Log.Information("Warpgate auth enabled: mutating endpoints require an authorized user from the {Header} header", Config.WarpgateAuth.UsernameHeader);
        }
        else
        {
            builder.AddTeleportAuth(Config.TeleportAuth);
            if (Config.TeleportAuth.Enabled)
                Log.Information("Teleport auth enabled: mutating endpoints require an authorized JWT (JWKS: {Jwks})", Config.TeleportAuth.JwksUrl);
            else
                Log.Warning("Teleport auth is DISABLED: mutating endpoints are gated by the API key only.");
        }

        builder.Services.AddSingleton<ICloudController, HetznerCloudController>(svc =>
        {
            ILogger<HetznerCloudController> logger = svc.GetRequiredService<ILogger<HetznerCloudController>>();
            return new HetznerCloudController(logger, Config.HetznerToken, Config.Sizes,
                Config.ProvisionScriptBaseUrl, Config.MetricUser, Config.MetricPassword);
        });
        
        builder.Services.AddSingleton<ICloudController, ProxmoxCloudController>(svc =>
        {
            ILogger<ProxmoxCloudController> logger = svc.GetRequiredService<ILogger<ProxmoxCloudController>>();
            return new ProxmoxCloudController(logger, Config.Sizes, Config.ProvisionScriptBaseUrl, Config.MetricUser, Config.MetricPassword, Config.PveHost, Config.PveUsername, Config.PvePassword, Config.PveTemplate, Config.MinVmId);
        });

        if (!string.IsNullOrWhiteSpace(Config.DigitalOceanToken))
        {
            builder.Services.AddSingleton<ICloudController, DigitalOceanCloudController>(svc =>
            {
                var logger = svc.GetRequiredService<ILogger<DigitalOceanCloudController>>();
                return new DigitalOceanCloudController(logger, Config.DigitalOceanToken,
                    Config.Sizes, Config.ProvisionScriptBaseUrl, Config.MetricUser,
                    Config.MetricPassword, Config.DigitalOceanRegions, Config.DigitalOceanTag,
                    Config.DigitalOceanVpcUuid, Config.DigitalOceanDefaultImage);
            });
        }

        if (!string.IsNullOrWhiteSpace(Config.AwsRegion))
        {
            builder.Services.AddSingleton<ICloudController, AwsCloudController>(svc =>
            {
                var logger = svc.GetRequiredService<ILogger<AwsCloudController>>();
                return new AwsCloudController(logger, Config.AwsRegion, Config.AwsAccessKeyId,
                    Config.AwsSecretAccessKey, Config.Sizes, Config.ProvisionScriptBaseUrl, Config.MetricUser,
                    Config.MetricPassword, Config.AwsSubnetIds, Config.AwsSecurityGroupIds, Config.AwsKeyName,
                    Config.AwsRootVolumeSizeGb, Config.AwsDefaultImage);
            });
        }

        WebApplication app = builder.Build();
        app.UseCors("AllowAll");

        // API key authentication for the dashboard/data API (/api/*).
        // Skipped for the health check and disabled entirely when no ApiKey is configured (local dev).
        app.UseWhen(
            ctx => ctx.Request.Path.StartsWithSegments("/api")
                   && !ctx.Request.Path.StartsWithSegments("/api/health"),
            branch => branch.Use(async (ctx, next) =>
            {
                string configuredKey = Program.Config.ApiKey;
                if (!string.IsNullOrEmpty(configuredKey))
                {
                    string provided = ctx.Request.Headers["X-API-KEY"].FirstOrDefault();
                    bool ok = !string.IsNullOrEmpty(provided) && CryptographicOperations.FixedTimeEquals(
                        Encoding.UTF8.GetBytes(provided), Encoding.UTF8.GetBytes(configuredKey));
                    if (!ok)
                    {
                        ctx.Response.StatusCode = StatusCodes.Status401Unauthorized;
                        await ctx.Response.WriteAsJsonAsync(new { message = "Invalid or missing API key" });
                        return;
                    }
                }

                await next();
            }));

        app.UseAuthentication();
        app.UseAuthorization();

        app.MapPost("/add-runner", AddRunnerManuallyHandler);
        app.MapPost("/runner-state", RunnerStateReportHandler);
        app.MapPost("/github-webhook", GithubWebhookHandler);
        app.MapControllers();
        foreach (var url in app.Urls)
        {
            Console.WriteLine($"Listening on {url}.");
        }

        Console.WriteLine("Start listening..."); 
        app.Run(Program.Config.ListenUrl);
    }

    static string GetFileHash(string filePath)
    {
        using SHA256 sha256 = SHA256.Create();
        // Open the file and compute the hash
        using FileStream stream = File.OpenRead(filePath);
        byte[] hashBytes = sha256.ComputeHash(stream);
        // Convert the byte array to a hexadecimal string
        StringBuilder hashString = new StringBuilder();
        foreach (byte b in hashBytes)
        {
            hashString.Append(b.ToString("x2")); // Convert each byte to a hexadecimal string
        }
        return hashString.ToString();
    }

    
    public static bool LoadConfiguration()
    {
        string configPath = Environment.GetEnvironmentVariable("CONFIG_FILE");
        if (string.IsNullOrEmpty(configPath))
        {
            string configDir = Environment.GetEnvironmentVariable("CONFIG_DIR") ??
                               Directory.CreateTempSubdirectory().FullName;
            configPath = Path.Combine(configDir, "config.json");
        }
        if (!File.Exists(configPath))
        {
            Log.Error($"Unable to locate config file at {configPath}");
            return false;
        }
        string configFileHash = GetFileHash(configPath);

        if (configFileHash == LoadedConfigHash)
        {
            return true;
        }
        
        Log.Information("Loading/refreshing configuration...");

        string configJson = File.ReadAllText(configPath);
        JsonSerializerOptions options = new()
        {
            Converters =
            {
                new JsonStringEnumConverter(JsonNamingPolicy.CamelCase)
            }
        };
        var loadedConfig = JsonSerializer.Deserialize<AutoScalerConfiguration>(configJson, options) ??
                 throw new Exception("Unable to parse configuration");

        if (string.IsNullOrWhiteSpace(loadedConfig.HetznerToken))
        {
            Log.Error($"Hetzner cloud token not set in {configPath}");
            return false;
        }
 
        Config = loadedConfig;
        LoadedConfigHash = configFileHash;
        Log.Information($"Loaded {Config.TargetConfigs.Count} targets and {Config.Sizes.Count} sizes.");
        return true;
    }

    public static string LoadedConfigHash { get; set; }

    private static async Task<IResult> GithubWebhookHandler(HttpRequest request, [FromServices] HetznerCloudController cloud, [FromServices] ILogger<Program> logger, [FromServices] RunnerQueue poolMgr)
    {
        using var activity = OrchestratorActivitySource.StartActivity("webhook.github");

        // Verify webhook HMAC
        request.EnableBuffering();
        string requestBody;
        using (var reader = new StreamReader(request.Body, leaveOpen: true))
        {
            requestBody = await reader.ReadToEndAsync();
            request.Body.Position = 0;
        }

        if (!string.IsNullOrEmpty(Config.WebhookSecret))
        {
            string signatureHeader = request.Headers["X-Hub-Signature-256"].FirstOrDefault();
            try
            {
                GithubSignatureVerifier.VerifySignature(requestBody, Config.WebhookSecret, signatureHeader);
            }
            catch (Exception ex)
            {
                activity?.SetStatus(ActivityStatusCode.Error, ex.Message);
                logger.LogWarning($"Webhook signature verification failed: {ex.Message}");
                return Results.StatusCode(401);
            }
        }

        // Read webhook from github
        JsonDocument json = JsonDocument.Parse(requestBody);

        string action;
        if (json.RootElement.TryGetProperty("action", out JsonElement actionJson))
        {
            action = actionJson.GetString() ?? string.Empty;
        }
        else
        {
            logger.LogDebug("No Action found. rejecting.");
            return Results.StatusCode(201);
        }

        activity?.SetTag("github.action", action);

        if (!json.RootElement.TryGetProperty("workflow_job", out JsonElement workflowJson))
        {
            logger.LogDebug("Received a non-workflowJob request. Ignoring.");
            return Results.StatusCode(201);
        }

        List<string> labels = workflowJson.GetProperty("labels")
            .EnumerateArray()
            .Select(x => x.GetString())
            .ToList();

        bool isSelfHosted = labels.Any(x => x.StartsWith($"self-hosted-{Program.Config.RunnerPrefix}"));

        if (!isSelfHosted)
        {
            //logger.LogDebug($"Received a non self-hosted request. Ignoring. Labels: {string.Join('|', labels)}");
            return Results.StatusCode(201);
        }

        long jobId = workflowJson.GetProperty("id").GetInt64();
        string repoNameRequest = json.RootElement.GetProperty("repository").GetProperty("full_name").GetString();
        string orgNameRequest = json.RootElement.GetProperty("organization").GetProperty("login").GetString();
        string jobUrl = workflowJson.GetProperty("url").GetString();

        // Needed to get properly cased names
        string orgName = Config.TargetConfigs.FirstOrDefault(x => x.Target == TargetType.Organization && x.Name.ToLower() == orgNameRequest.ToLower())?.Name ?? orgNameRequest;
        string repoName = Config.TargetConfigs.FirstOrDefault(x => x.Target == TargetType.Repository && x.Name.ToLower() == repoNameRequest.ToLower())?.Name ?? repoNameRequest;

        activity?.SetTag("github.job_id", jobId);
        activity?.SetTag("github.repo", repoName);
        activity?.SetTag("github.org", orgName);

        // Check if its an org or a repo
        if (String.IsNullOrEmpty(orgName))
        {
            logger.LogWarning($"Unable to retrieve organization. aborting.");
            return Results.StatusCode(201);
        }

        // Check if Org is configured
        bool isOrg = Config.TargetConfigs.Any(x => x.Name == orgName && x.Target == TargetType.Organization);
        bool isRepo = Config.TargetConfigs.Any(x => x.Name == repoName && x.Target == TargetType.Repository);

        if (!isOrg && !isRepo)
        {
            logger.LogWarning($"GitHub org {orgName} nor repo {repoName} is configured. Ignoring.");
            return Results.StatusCode(201);
        }


        await using var db = new ActionsRunnerContext();

        try
        {
            switch (action)
            {
                case "queued":
                    using (OrchestratorActivitySource.StartActivity("webhook.github.queued"))
                    {
                        await JobQueued(logger, repoName, labels, orgName, poolMgr, isRepo ? TargetType.Repository : TargetType.Organization, jobId, jobUrl);
                    }
                    break;
                case "in_progress":
                    using (OrchestratorActivitySource.StartActivity("webhook.github.in_progress"))
                    {
                        var dbWorkflow = await db.Jobs.FirstOrDefaultAsync(x => x.GithubJobId == jobId);
                        if (dbWorkflow == null)
                        {
                            logger.LogWarning("Processing job on manually created runner");
                            Job progressJob = new()
                            {
                                GithubJobId = jobId,
                                Repository = repoName,
                                Owner = isRepo ? repoName : orgName,
                                State = JobState.InProgress,
                                InProgressTime = DateTime.UtcNow,
                                JobUrl = jobUrl,
                                Orphan = true
                            };
                            await db.Jobs.AddAsync(progressJob);
                        }
                        else
                        {
                            dbWorkflow.State = JobState.InProgress;
                            dbWorkflow.QueueTime = DateTime.UtcNow;
                        }
                        await db.SaveChangesAsync();
                        await JobInProgress(workflowJson, logger, jobId, repoName, orgName);
                    }
                    break;
                case "completed":
                    using (OrchestratorActivitySource.StartActivity("webhook.github.completed"))
                    {
                        string conclusion = String.Empty;
                        if (json.RootElement.TryGetProperty("conclusion", out JsonElement conclusionJson))
                        {
                            conclusion = conclusionJson.GetString() ?? string.Empty;
                        }

                        var dbWorkflowComplete = await db.Jobs.FirstOrDefaultAsync(x => x.GithubJobId == jobId);
                        if (dbWorkflowComplete == null)
                        {
                            logger.LogWarning($"Completed webhook for unknown job {jobId} in {repoName}. Creating record.");
                            dbWorkflowComplete = new Job
                            {
                                GithubJobId = jobId,
                                Repository = repoName,
                                Owner = isRepo ? repoName : orgName,
                                State = JobState.Completed,
                                CompleteTime = DateTime.UtcNow,
                                Orphan = true
                            };
                            await db.Jobs.AddAsync(dbWorkflowComplete);
                            await db.SaveChangesAsync();
                            return Results.StatusCode(201);
                        }
                        dbWorkflowComplete.CompleteTime = DateTime.UtcNow;
                        bool wasCancelled = false;
                        switch (conclusion)
                        {
                            case "cancelled":
                                dbWorkflowComplete.State = JobState.Cancelled;
                                await db.SaveChangesAsync();
                                wasCancelled = true;
                                break;
                            default:
                                dbWorkflowComplete.State = JobState.Completed;
                                await db.SaveChangesAsync();
                                break;
                        }

                        await JobCompleted(logger, jobId, poolMgr, repoName, orgName, workflowJson, wasCancelled);
                    }
                    break;
                default:
                    logger.LogWarning("Unknown action. Ignoring");
                    break;
            }
        }
        catch (Exception ex)
        {
            activity?.SetStatus(ActivityStatusCode.Error, ex.Message);
            activity?.AddException(ex);
            // This should make the webhook as bad and the timer will redeliver it after a while
            Log.Error($"Failed to process {action} webhook: {ex.Message}");
            return Results.StatusCode(500);
        }

        // All was well
        return Results.StatusCode(201);
    }

    // AWS/DO assign the public IP after the create call returns - poll the CSP until it shows up
    private static async Task<string> WaitForPublicIp(ICloudController cloudController, long serverId, int timeoutSeconds = 60, int pollSeconds = 2)
    {
        var endTime = DateTime.UtcNow.AddSeconds(timeoutSeconds);
        while (DateTime.UtcNow < endTime)
        {
            var ip = await cloudController.GetRunnerIpAddress(serverId);
            if (!string.IsNullOrEmpty(ip) && ip != "0.0.0.0/0")
                return ip;
            await Task.Delay(TimeSpan.FromSeconds(pollSeconds));
        }
        return string.Empty;
    }

    private static async Task<IResult> RunnerStateReportHandler(HttpRequest request, [FromServices] IServiceProvider serviceProvider , [FromServices] ILogger<Program> logger, [FromServices] RunnerQueue runnerQueue, [FromQuery] string hostname, [FromQuery] string state)
    {
        await using var db = new ActionsRunnerContext();
        var runner = await db.Runners.Include(x => x.Lifecycle).Include(x => x.Job).FirstOrDefaultAsync(x => x.Hostname == hostname);
        if (runner == null)
        {
            logger.LogError("Runner {Hostname} not found in database during state report ({State}). Removing tracking entry.", hostname, state);
            runnerQueue.CreatedRunners.Remove(hostname, out _);
            return Results.StatusCode(404);
        }
        switch (state)
        {
            case "ok":
                // Remove from provisioning dict
                Log.Information($"Runner {hostname} finished provisioning.");
                MachineSuccessCount.Labels(runner.Arch, runner.Size).Inc();
                runnerQueue.CreatedRunners.Remove(hostname, out _);
                runner.Lifecycle.Add(new()
                {
                    Event = "Runner finished provisioning",
                    Status = RunnerStatus.Provisioned,
                    EventTimeUtc = DateTime.UtcNow
                });
                
                // Update the runner's IP address if it's still missing or a dummy value
                // (e.g. AWS/DO public IPs that were not yet assigned when the machine was created)
                if (string.IsNullOrEmpty(runner.IPv4) || runner.IPv4 == "0.0.0.0/0")
                {
                    try
                    {
                        var cloudController = serviceProvider.GetServices<ICloudController>()
                            .FirstOrDefault(x => x.CloudIdentifier == runner.Cloud);
                        if (cloudController != null)
                        {
                            var actualIpAddress = await WaitForPublicIp(cloudController, runner.CloudServerId);
                            if (!string.IsNullOrEmpty(actualIpAddress))
                            {
                                runner.IPv4 = actualIpAddress;
                                logger.LogInformation($"Updated IP for runner {runner.Hostname} ({runner.Cloud}) to {actualIpAddress}");
                                await db.SaveChangesAsync();
                            }
                        }
                    }
                    catch (Exception ex)
                    {
                        // Log the error but don't fail the provision request
                        SentrySdk.CaptureException(ex);
                        Console.WriteLine($"Failed to update IP for runner {runner.RunnerId}: {ex.Message}");
                    }
                }
                
                
                await db.SaveChangesAsync();
                break;
            case "error":
                Log.Warning($"Runner {hostname} failed provisioning.");
                MachineFailedCount.Labels(runner.Arch, runner.Size).Inc();
                if (request.Form.Files.Count > 0)
                {
                    // Read the log file into a string
                    using var reader = new StreamReader(request.Form.Files[0].OpenReadStream());
                    string fileContent = await reader.ReadToEndAsync();
                    Log.Information($"LOGS FROM {hostname}\n\n{fileContent}");
                }

                runner.Lifecycle.Add(new()
                {
                    Event = "Runner failed provisioning",
                    Status = RunnerStatus.Failure,
                    EventTimeUtc = DateTime.UtcNow
                });
                // Get runner specs

                // Queue creation of a new runner
                if (!runnerQueue.CreatedRunners.Remove(hostname, out CreateRunnerTask runnerSpec))
                {
                    logger.LogError($"Unable to get previous settings for {hostname}");
                }
                else
                {
                    Runner newRunner = new()
                    {
                        Size = runner.Size,
                        // Cloud is only known once the PoolManager picked a provider
                        Cloud = "ukn",
                        Hostname = "Unknown",
                        Profile = runner.Profile,
                        Lifecycle =
                        [
                            new RunnerLifecycle
                            {
                                EventTimeUtc = DateTime.UtcNow,
                                Status = RunnerStatus.CreationQueued,
                                Event = $"Created as replacement for failed {hostname}"
                            }
                        ],
                        IsOnline = false,
                        Arch = runner.Arch,
                        IPv4 = string.Empty,
                        IsCustom = runner.IsCustom,
                        Owner = runner.Owner
                    };
                    await db.Runners.AddAsync(newRunner);
                    await db.SaveChangesAsync();

                    runnerSpec.RunnerDbId = newRunner.RunnerId;

                    logger.LogInformation($"Re-creating runner of type [{runner.Size}, {runner.Arch}] for {runner.Job.Repository}");
                    runnerQueue.CreateTasks.Enqueue(runnerSpec);

                    // Queue deletion of the failed runner - unless it is held for inspection
                    if (runner.IsHeld)
                    {
                        logger.LogWarning($"Runner {runner.Hostname} failed provisioning but is held for inspection - not queueing deletion");
                    }
                    else
                    {
                        runnerQueue.DeleteTasks.Enqueue(new DeleteRunnerTask { ServerId = runner.CloudServerId });
                    }
                    await db.SaveChangesAsync();
                }

                break;
            default:
                logger.LogWarning("Runner {Hostname} reported unknown state '{State}'. Removing tracking entry.", hostname, state);
                runnerQueue.CreatedRunners.Remove(hostname, out _);
                break;
        }

        return Results.StatusCode(201);
    }

    private static async Task<IResult> AddRunnerManuallyHandler(HttpRequest request, [FromServices] HetznerCloudController cloud, [FromServices] ILogger<Program> logger, [FromServices] RunnerQueue runnerQueue, [FromServices] RunnerQueue poolMgr)
    {
        // Read webhook from github
        JsonDocument json = await request.ReadFromJsonAsync<JsonDocument>();

        string repoNameRequest = json.RootElement.GetProperty("repo").GetString();
        string orgNameRequest = json.RootElement.GetProperty("org").GetString();
        List<string> labels = json.RootElement.GetProperty("labels")
            .EnumerateArray()
            .Select(x => x.GetString())
            .ToList();

        // Needed to get properly cased names
        string orgName = Config.TargetConfigs.FirstOrDefault(x => x.Target == TargetType.Organization && x.Name.ToLower() == orgNameRequest.ToLower())?.Name ?? orgNameRequest;
        string repoName = Config.TargetConfigs.FirstOrDefault(x => x.Target == TargetType.Repository && x.Name.ToLower() == repoNameRequest.ToLower())?.Name ?? repoNameRequest;


        // Check if its an org or a repo
        if (String.IsNullOrEmpty(orgName))
        {
            logger.LogWarning($"Unable to retrieve organization. aborting.");
            return Results.StatusCode(400);
        }

        // Check if Org is configured
        bool isOrg = Config.TargetConfigs.Any(x => x.Name == orgName && x.Target == TargetType.Organization);
        bool isRepo = Config.TargetConfigs.Any(x => x.Name == repoName && x.Target == TargetType.Repository);

        if (!isOrg && !isRepo)
        {
            logger.LogWarning($"GitHub org {orgName} nor repo {repoName} is configured. Ignoring.");
            return Results.StatusCode(400);
        }

        try
        {
            await JobQueued(logger, repoName, labels, orgName, poolMgr, isRepo ? TargetType.Repository : TargetType.Organization, -1, null);
        }
        catch (Exception ex)
        {
            // This should make the webhook as bad and the timer will redeliver it after a while
            Log.Error($"Failed to process manual trigger: {ex.Message}");
            return Results.StatusCode(500);
        }

        // All was well 
        return Results.StatusCode(201);
    }

    private static async Task JobCompleted(ILogger<Program> logger, long jobId, RunnerQueue poolMgr, string repoName, string orgName, JsonElement workflowJson, bool wasCancelled)
    {
        await using var db = new ActionsRunnerContext();
        var job = await db.Jobs
            .Include(job => job.Runner)
            .ThenInclude(runner => runner.Lifecycle)
            .FirstOrDefaultAsync(x => x.GithubJobId == jobId);

        if (job == null)
        {
            logger.LogWarning($"No job in database with ID: {jobId} on repo {repoName}");
            await db.Jobs.AddAsync(new Job
            {
                GithubJobId = jobId,
                Repository = repoName,
                Owner = orgName,
                State = JobState.Completed,
                Orphan = true,
                CompleteTime = DateTime.UtcNow
            });
            await db.SaveChangesAsync();
            return;
        }
        
        logger.LogInformation(
            $"Workflow Job {jobId} in repo {repoName} has completed. Queuing deletion of VM associated with Job.");

        Runner jobRunner = null;
        if (job.Runner == null)
        {
            // Retroactivly assign runner to job
            string runnerName = workflowJson.GetProperty("runner_name").GetString();
            logger.LogError($"No VM on record for JobID: {jobId}. Trying to re-link to {runnerName}.");
            jobRunner = await db.LinkJobToRunner(jobId, runnerName);

            if (wasCancelled && jobRunner == null)
            {
                // job was cancelled before a runner was picked by the job
                // it's either still in the creation queue or it is still provisioning

                // Get arch from the requested size
                string arch = Config.Sizes.FirstOrDefault(x => x.Name == job.RequestedSize)?.Arch ?? "x64";

                // Build the cancellation key
                var key = (job.Owner, job.Repository, job.RequestedSize, job.RequestedProfile, arch);

                // Increment the cancelled runners counter
                int count = poolMgr.CancelledRunners.AddOrUpdate(key, 1, (k, currentCount) => currentCount + 1);

                // Increment Prometheus counter
                CancelledJobCount.Labels(job.Owner, job.Repository, job.RequestedSize, job.RequestedProfile ?? "default", arch).Inc();

                logger.LogInformation($"Registered cancelled job {jobId} for {job.Owner}/{job.Repository} size={job.RequestedSize} profile={job.RequestedProfile} arch={arch}. Counter now at {count}.");

                return;
            }
            else if (jobRunner == null)
            {
                logger.LogError("Unable to link runner. aborting");
                return;
            }
        }
        else
        {
            jobRunner = job.Runner;
        }
        
        // record event in DB
        if (jobRunner.IsHeld)
        {
            // Runner is held for manual inspection: the job is done, but the VM is kept
            // for the operator to investigate. It will only go away on explicit deletion.
            jobRunner.Lifecycle.Add(new()
            {
                Status = RunnerStatus.Provisioned,
                EventTimeUtc = DateTime.UtcNow,
                Event = $"Workflow Job {jobId} in repo {repoName} has completed. Runner is held for inspection - keeping VM."
            });
        }
        else if (jobRunner.Lifecycle.Any(x => x.Status == RunnerStatus.DeletionQueued))
        {
            jobRunner.Lifecycle.Add(new()
            {
                Status = RunnerStatus.DeletionQueued,
                EventTimeUtc = DateTime.UtcNow,
                Event = $"Workflow Job {jobId} in repo {repoName} has completed. Deletion already queued."
            });
            jobRunner.IsOnline = false;
            
        }
        else
        {
            jobRunner.Lifecycle.Add(new()
            {
                Status = RunnerStatus.DeletionQueued,
                EventTimeUtc = DateTime.UtcNow,
                Event = $"Workflow Job {jobId} in repo {repoName} has completed. Deletion queued."
            });
            jobRunner.IsOnline = false;
            // Sent to pool manager to delete
            poolMgr.DeleteTasks.Enqueue(new DeleteRunnerTask
            {
                ServerId = jobRunner.CloudServerId,
                RunnerDbId = jobRunner.RunnerId
            });
        }

        await db.SaveChangesAsync();
        
        ProcessedJobCount.Labels(job.Owner, jobRunner.Size).Inc();

        double secondsAlive = (DateTime.UtcNow - jobRunner.CreationQueuedTime).TotalSeconds;
        TotalMachineTime.Labels(job.Owner, jobRunner.Size).Inc(secondsAlive);
        
    }

    private static async Task JobInProgress(JsonElement workflowJson, ILogger<Program> logger, long jobId,
        string repoName, string orgName)
    {
        string runnerName = workflowJson.GetProperty("runner_name").GetString();
        logger.LogInformation($"Workflow Job {jobId} in repo {repoName} now in progress on {runnerName}");
        
        // Make the connection between the job and the runner in the DB
        await using var db = new ActionsRunnerContext();
        Runner runner = await db.LinkJobToRunner(jobId, runnerName);

        if (runner == null)
        {
            logger.LogWarning($"Unable to link job {jobId} to runner {runnerName}");
            return;
        }

        // Metrics
        PickedJobCount.Labels(orgName, runner.Size).Inc();

    }

    private static async Task JobQueued(ILogger<Program> logger, string repoName, List<string> labels, string orgName, RunnerQueue poolMgr, TargetType targetType, long jobId, string jobUrl)
    {
        logger.LogInformation($"New Workflow Job was queued for {repoName}. Queuing VM creation to replenish pool...");
        
        string size = string.Empty;
        string arch = string.Empty;
        string profileName = "default";
        bool isCustom = false;

        // Check if this is a custom run
        if (labels.Contains($"self-hosted-{Config.RunnerPrefix}-custom"))
        {
            isCustom = true;
            // Check for a profile label
            string profileLabel = labels.FirstOrDefault(x => x.StartsWith("profile-"));
            if (string.IsNullOrEmpty(profileLabel))
            {
                logger.LogError("No profile label given for custom mode. Ignoring.");
                return;
            }

            profileName = profileLabel.Split('-')[1];

        }
        foreach (MachineSize csize in Config.Sizes)
        {
            if (labels.Contains(csize.Name) || labels.Contains($"self-hosted-{Config.RunnerPrefix}-{csize.Name}"))
            {
                size = csize.Name;
                arch = csize.Arch;
                break;
            }
        }

        if (string.IsNullOrEmpty(size))
        {
            logger.LogWarning($"No runner size specified for workflow in {repoName}. Ignoring.");
            return;
        }

        // Create a new runner
        string githubToken = targetType switch
        {
            TargetType.Organization => Config.TargetConfigs.FirstOrDefault(x => x.Name == orgName && x.Target == TargetType.Organization)?.GitHubToken,
            TargetType.Repository => Config.TargetConfigs.FirstOrDefault(x => x.Name == repoName && x.Target == TargetType.Repository)?.GitHubToken,
            _ => throw new ArgumentOutOfRangeException(nameof(targetType), targetType, null)
        }; 
       
        if (String.IsNullOrEmpty(githubToken))
        {
            switch (targetType)
            {
                case TargetType.Organization:
                    logger.LogError($"Unknown organization: {orgName} - check setup. aborting.");
                    break;
                case TargetType.Repository:
                    logger.LogError($"Unknown repository: {repoName} - check setup. aborting.");
                    break;
            }

            return;
        }

        string owner = targetType switch
        {
            TargetType.Organization => orgName,
            TargetType.Repository => repoName,
            _ => throw new ArgumentOutOfRangeException(nameof(targetType), targetType, null)
        };

        // Get the target configuration to check quota
        var targetConfig = targetType switch
        {
            TargetType.Organization => Config.TargetConfigs.FirstOrDefault(x => x.Name == orgName && x.Target == TargetType.Organization),
            TargetType.Repository => Config.TargetConfigs.FirstOrDefault(x => x.Name == repoName && x.Target == TargetType.Repository),
            _ => throw new ArgumentOutOfRangeException(nameof(targetType), targetType, null)
        };

        // Record runner to database
        await using var db = new ActionsRunnerContext();

        // Check if quota is reached before creating runner
        if (targetConfig?.RunnerQuota.HasValue == true)
        {
            // Count all runners that are actively consuming resources (not deleted/failed/cancelled)
            // This includes: CreationQueued, Created, Provisioned, Processing (states 1-4)
            var runners = await db.Runners
                .Include(x => x.Lifecycle)
                .Where(x => x.Owner == owner)
                .ToListAsync();

            int currentRunnerCount = runners.Count(x => x.LastState < RunnerStatus.DeletionQueued);

            if (currentRunnerCount >= targetConfig.RunnerQuota.Value)
            {
                logger.LogWarning($"Runner quota reached for {owner}: {currentRunnerCount}/{targetConfig.RunnerQuota.Value} (includes queued/provisioning runners). Job will be throttled.");

                // Record the job in the database as Throttled
                if (jobId > 0)
                {
                    Job throttledJob = new()
                    {
                        GithubJobId = jobId,
                        Repository = repoName,
                        Owner = owner,
                        State = JobState.Throttled,
                        QueueTime = DateTime.UtcNow,
                        JobUrl = jobUrl,
                        Orphan = false,
                        RequestedProfile = profileName,
                        RequestedSize = size
                    };
                    await db.Jobs.AddAsync(throttledJob);
                    await db.SaveChangesAsync();
                }

                return;
            }
        }

        if (jobId > 0)
        {
            Job queuedJob = new()
            {
                GithubJobId = jobId,
                Repository = repoName,
                Owner = owner,
                State = JobState.Queued,
                QueueTime = DateTime.UtcNow,
                JobUrl = jobUrl,
                Orphan = false,
                RequestedProfile = profileName,
                RequestedSize = size
            };
            await db.Jobs.AddAsync(queuedJob);
        }

        Runner newRunner = new()
        {
            Size = size,
            // Cloud is only known once the PoolManager picked a provider
            Cloud = "ukn",
            Hostname = "Unknown",
            Profile = profileName,
            Lifecycle =
            [
                new RunnerLifecycle
                {
                    EventTimeUtc = DateTime.UtcNow,
                    Status = RunnerStatus.CreationQueued,
                    Event = "Created as queued job runner"
                }
            ],
            IsOnline = false,
            Arch = arch,
            IPv4 = string.Empty,
            IsCustom = isCustom,
            Owner = owner
            
        };
        await db.Runners.AddAsync(newRunner);
        await db.SaveChangesAsync();
        
        poolMgr.CreateTasks.Enqueue(new CreateRunnerTask
        {
            RepoName = repoName,
            TargetType = targetType,
            RunnerDbId = newRunner.RunnerId
        });

        QueuedJobCount.Labels(orgName, size).Inc();
    }
}