using Dapper;
using KArtSell.BuildingBlocks.Data;

namespace KArtSell.BuildingBlocks.Reliability;

public sealed class DapperJobRunRepository(IDbConnectionFactory connectionFactory) : IJobRunRepository
{
    private const string InsertSql = """
        insert into building_blocks.job_run
        (job_run_id, job_type, job_version, scope_key, idempotency_key, watermark,
         app_version, model_version, config_version, data_version, contract_version,
         status, input_hash, started_at, heartbeat_at, trace_id)
        values
        (@JobRunId, @JobType, @JobVersion, @ScopeKey, @IdempotencyKey, @Watermark,
         @CodeSha, @ModelVersion, @ConfigVersion, @DatasetId, @ContractVersion,
         'Running', @InputHash, @RequestedAt, @RequestedAt, @TraceId)
        on conflict (idempotency_key) do nothing;
        """;

    private const string HeartbeatSql = """
        update building_blocks.job_run
        set heartbeat_at = @At
        where job_run_id = @JobRunId and status = 'Running';
        """;

    private const string CompleteSql = """
        update building_blocks.job_run
        set status = @Status,
            output_hash = @OutputHash,
            watermark = coalesce(@Watermark, watermark),
            error_code = @ErrorCode,
            finished_at = @FinishedAt,
            heartbeat_at = @FinishedAt
        where job_run_id = @JobRunId and status = 'Running';
        """;

    public async Task<bool> TryStartAsync(JobRunStart start, CancellationToken cancellationToken)
    {
        start.VersionSet.EnsureValid();
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        var affected = await connection.ExecuteAsync(new CommandDefinition(
            InsertSql,
            new
            {
                start.JobRunId,
                start.JobType,
                start.JobVersion,
                start.ScopeKey,
                start.IdempotencyKey,
                start.Watermark,
                start.VersionSet.CodeSha,
                start.VersionSet.ModelVersion,
                start.VersionSet.ConfigVersion,
                start.VersionSet.DatasetId,
                start.VersionSet.ContractVersion,
                start.InputHash,
                start.TraceId,
                start.RequestedAt
            },
            cancellationToken: cancellationToken));
        return affected == 1;
    }

    public async Task HeartbeatAsync(Guid jobRunId, DateTimeOffset at, CancellationToken cancellationToken)
    {
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        await connection.ExecuteAsync(new CommandDefinition(
            HeartbeatSql,
            new { JobRunId = jobRunId, At = at },
            cancellationToken: cancellationToken));
    }

    public async Task CompleteAsync(JobRunCompletion completion, CancellationToken cancellationToken)
    {
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        var affected = await connection.ExecuteAsync(new CommandDefinition(
            CompleteSql,
            new
            {
                completion.JobRunId,
                Status = completion.Status.ToString(),
                completion.OutputHash,
                completion.Watermark,
                completion.ErrorCode,
                completion.FinishedAt
            },
            cancellationToken: cancellationToken));
        if (affected != 1)
        {
            throw new InvalidOperationException("Job run is absent or no longer Running.");
        }
    }
}
