using System;
using System.Collections.Immutable;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using MongoDB.Driver;
using StellaOps.Scheduler.Models;
using StellaOps.Scheduler.Storage.Mongo.Repositories;
using StellaOps.Scheduler.Worker.Options;
using StellaOps.Scheduler.Worker.Observability;
using StellaOps.Scheduler.Worker.Policy;
using Xunit;

namespace StellaOps.Scheduler.Worker.Tests;

public sealed class PolicyRunExecutionServiceTests
{
    private static readonly SchedulerWorkerOptions WorkerOptions = new()
    {
        Policy =
        {
            Dispatch =
            {
                LeaseOwner = "test-dispatch",
                BatchSize = 1,
                LeaseDuration = TimeSpan.FromMinutes(1),
                IdleDelay = TimeSpan.FromMilliseconds(10),
                MaxAttempts = 2,
                RetryBackoff = TimeSpan.FromSeconds(30)
            },
            Api =
            {
                BaseAddress = new Uri("https://policy.example.com"),
                RunsPath = "/api/policy/policies/{policyId}/runs",
                SimulatePath = "/api/policy/policies/{policyId}/simulate"
            }
        }
    };

    [Fact]
    public async Task ExecuteAsync_CancelsJob_WhenCancellationRequested()
    {
        var repository = new RecordingPolicyRunJobRepository();
        var client = new StubPolicyRunClient();
        var options = Microsoft.Extensions.Options.Options.Create(CloneOptions());
        var timeProvider = new TestTimeProvider(DateTimeOffset.Parse("2025-10-28T10:00:00Z"));
        using var metrics = new SchedulerWorkerMetrics();
        var targeting = new StubPolicyRunTargetingService
        {
            OnEnsureTargets = job => PolicyRunTargetingResult.Unchanged(job)
        };
        var service = new PolicyRunExecutionService(repository, client, options, timeProvider, metrics, targeting, NullLogger<PolicyRunExecutionService>.Instance);

        var job = CreateJob(status: PolicyRunJobStatus.Dispatching) with
        {
            CancellationRequested = true,
            LeaseOwner = "test-dispatch",
            LeaseExpiresAt = timeProvider.GetUtcNow().AddMinutes(1)
        };

        var result = await service.ExecuteAsync(job, CancellationToken.None);

        Assert.Equal(PolicyRunExecutionResultType.Cancelled, result.Type);
        Assert.Equal(PolicyRunJobStatus.Cancelled, result.UpdatedJob.Status);
        Assert.True(repository.ReplaceCalled);
        Assert.Equal("test-dispatch", repository.ExpectedLeaseOwner);
    }

    [Fact]
    public async Task ExecuteAsync_SubmitsJob_OnSuccess()
    {
        var repository = new RecordingPolicyRunJobRepository();
        var client = new StubPolicyRunClient
        {
            Result = PolicyRunSubmissionResult.Succeeded("run:P-7:2025", DateTimeOffset.Parse("2025-10-28T10:01:00Z"))
        };
        var options = Microsoft.Extensions.Options.Options.Create(CloneOptions());
        var timeProvider = new TestTimeProvider(DateTimeOffset.Parse("2025-10-28T10:00:00Z"));
        using var metrics = new SchedulerWorkerMetrics();
        var targeting = new StubPolicyRunTargetingService
        {
            OnEnsureTargets = job => PolicyRunTargetingResult.Unchanged(job)
        };
        var service = new PolicyRunExecutionService(repository, client, options, timeProvider, metrics, targeting, NullLogger<PolicyRunExecutionService>.Instance);

        var job = CreateJob(status: PolicyRunJobStatus.Dispatching) with
        {
            LeaseOwner = "test-dispatch",
            LeaseExpiresAt = timeProvider.GetUtcNow().AddMinutes(1)
        };

        var result = await service.ExecuteAsync(job, CancellationToken.None);

        Assert.Equal(PolicyRunExecutionResultType.Submitted, result.Type);
        Assert.Equal(PolicyRunJobStatus.Submitted, result.UpdatedJob.Status);
        Assert.Equal("run:P-7:2025", result.UpdatedJob.RunId);
        Assert.Equal(job.AttemptCount + 1, result.UpdatedJob.AttemptCount);
        Assert.Null(result.UpdatedJob.LastError);
        Assert.True(repository.ReplaceCalled);
    }

    [Fact]
    public async Task ExecuteAsync_RetriesJob_OnFailure()
    {
        var repository = new RecordingPolicyRunJobRepository();
        var client = new StubPolicyRunClient
        {
            Result = PolicyRunSubmissionResult.Failed("timeout")
        };
        var options = Microsoft.Extensions.Options.Options.Create(CloneOptions());
        var timeProvider = new TestTimeProvider(DateTimeOffset.Parse("2025-10-28T10:00:00Z"));
        using var metrics = new SchedulerWorkerMetrics();
        var targeting = new StubPolicyRunTargetingService
        {
            OnEnsureTargets = job => PolicyRunTargetingResult.Unchanged(job)
        };
        var service = new PolicyRunExecutionService(repository, client, options, timeProvider, metrics, targeting, NullLogger<PolicyRunExecutionService>.Instance);

        var job = CreateJob(status: PolicyRunJobStatus.Dispatching) with
        {
            LeaseOwner = "test-dispatch",
            LeaseExpiresAt = timeProvider.GetUtcNow().AddMinutes(1)
        };

        var result = await service.ExecuteAsync(job, CancellationToken.None);

        Assert.Equal(PolicyRunExecutionResultType.Retrying, result.Type);
        Assert.Equal(PolicyRunJobStatus.Pending, result.UpdatedJob.Status);
        Assert.Equal(job.AttemptCount + 1, result.UpdatedJob.AttemptCount);
        Assert.Equal("timeout", result.UpdatedJob.LastError);
        Assert.True(result.UpdatedJob.AvailableAt > job.AvailableAt);
    }

    [Fact]
    public async Task ExecuteAsync_MarksJobFailed_WhenAttemptsExceeded()
    {
        var repository = new RecordingPolicyRunJobRepository();
        var client = new StubPolicyRunClient
        {
            Result = PolicyRunSubmissionResult.Failed("bad request")
        };
        var optionsValue = CloneOptions();
        optionsValue.Policy.Dispatch.MaxAttempts = 1;
        var options = Microsoft.Extensions.Options.Options.Create(optionsValue);
        var timeProvider = new TestTimeProvider(DateTimeOffset.Parse("2025-10-28T10:00:00Z"));
        using var metrics = new SchedulerWorkerMetrics();
        var targeting = new StubPolicyRunTargetingService
        {
            OnEnsureTargets = job => PolicyRunTargetingResult.Unchanged(job)
        };
        var service = new PolicyRunExecutionService(repository, client, options, timeProvider, metrics, targeting, NullLogger<PolicyRunExecutionService>.Instance);

        var job = CreateJob(status: PolicyRunJobStatus.Dispatching, attemptCount: 0) with
        {
            LeaseOwner = "test-dispatch",
            LeaseExpiresAt = timeProvider.GetUtcNow().AddMinutes(1)
        };

        var result = await service.ExecuteAsync(job, CancellationToken.None);

        Assert.Equal(PolicyRunExecutionResultType.Failed, result.Type);
        Assert.Equal(PolicyRunJobStatus.Failed, result.UpdatedJob.Status);
        Assert.Equal("bad request", result.UpdatedJob.LastError);
    }

    [Fact]
    public async Task ExecuteAsync_NoWork_CompletesJob()
    {
        var repository = new RecordingPolicyRunJobRepository();
        var client = new StubPolicyRunClient();
        var options = Microsoft.Extensions.Options.Options.Create(CloneOptions());
        var timeProvider = new TestTimeProvider(DateTimeOffset.Parse("2025-10-28T10:00:00Z"));
        using var metrics = new SchedulerWorkerMetrics();
        var targeting = new StubPolicyRunTargetingService
        {
            OnEnsureTargets = job => PolicyRunTargetingResult.NoWork(job, "empty")
        };
        var service = new PolicyRunExecutionService(repository, client, options, timeProvider, metrics, targeting, NullLogger<PolicyRunExecutionService>.Instance);

        var job = CreateJob(status: PolicyRunJobStatus.Dispatching, inputs: PolicyRunInputs.Empty) with
        {
            LeaseOwner = "test-dispatch",
            LeaseExpiresAt = timeProvider.GetUtcNow().AddMinutes(1)
        };

        var result = await service.ExecuteAsync(job, CancellationToken.None);

        Assert.Equal(PolicyRunExecutionResultType.NoOp, result.Type);
        Assert.Equal(PolicyRunJobStatus.Completed, result.UpdatedJob.Status);
        Assert.True(repository.ReplaceCalled);
        Assert.Equal("test-dispatch", repository.ExpectedLeaseOwner);
    }

    private static PolicyRunJob CreateJob(PolicyRunJobStatus status, int attemptCount = 0, PolicyRunInputs? inputs = null)
    {
        var resolvedInputs = inputs ?? new PolicyRunInputs(sbomSet: new[] { "sbom:S-42" }, captureExplain: true);
        var metadata = ImmutableSortedDictionary.Create<string, string>(StringComparer.Ordinal);
        return new PolicyRunJob(
            SchemaVersion: SchedulerSchemaVersions.PolicyRunJob,
            Id: "job_1",
            TenantId: "tenant-alpha",
            PolicyId: "P-7",
            PolicyVersion: 4,
            Mode: PolicyRunMode.Incremental,
            Priority: PolicyRunPriority.Normal,
            PriorityRank: -1,
            RunId: "run:P-7:2025",
            RequestedBy: "user:cli",
            CorrelationId: "corr-1",
            Metadata: metadata,
            Inputs: resolvedInputs,
            QueuedAt: DateTimeOffset.Parse("2025-10-28T09:59:00Z"),
            Status: status,
            AttemptCount: attemptCount,
            LastAttemptAt: null,
            LastError: null,
            CreatedAt: DateTimeOffset.Parse("2025-10-28T09:58:00Z"),
            UpdatedAt: DateTimeOffset.Parse("2025-10-28T09:58:00Z"),
            AvailableAt: DateTimeOffset.Parse("2025-10-28T09:59:00Z"),
            SubmittedAt: null,
            CompletedAt: null,
            LeaseOwner: null,
            LeaseExpiresAt: null,
            CancellationRequested: false,
            CancellationRequestedAt: null,
            CancellationReason: null,
            CancelledAt: null);
    }

    private static SchedulerWorkerOptions CloneOptions()
    {
        return new SchedulerWorkerOptions
        {
            Policy = new SchedulerWorkerOptions.PolicyOptions
            {
                Enabled = WorkerOptions.Policy.Enabled,
                Dispatch = new SchedulerWorkerOptions.PolicyOptions.DispatchOptions
                {
                    LeaseOwner = WorkerOptions.Policy.Dispatch.LeaseOwner,
                    BatchSize = WorkerOptions.Policy.Dispatch.BatchSize,
                    LeaseDuration = WorkerOptions.Policy.Dispatch.LeaseDuration,
                    IdleDelay = WorkerOptions.Policy.Dispatch.IdleDelay,
                    MaxAttempts = WorkerOptions.Policy.Dispatch.MaxAttempts,
                    RetryBackoff = WorkerOptions.Policy.Dispatch.RetryBackoff
                },
                Api = new SchedulerWorkerOptions.PolicyOptions.ApiOptions
                {
                    BaseAddress = WorkerOptions.Policy.Api.BaseAddress,
                    RunsPath = WorkerOptions.Policy.Api.RunsPath,
                    SimulatePath = WorkerOptions.Policy.Api.SimulatePath,
                    TenantHeader = WorkerOptions.Policy.Api.TenantHeader,
                    IdempotencyHeader = WorkerOptions.Policy.Api.IdempotencyHeader,
                    RequestTimeout = WorkerOptions.Policy.Api.RequestTimeout
                },
                Targeting = new SchedulerWorkerOptions.PolicyOptions.TargetingOptions
                {
                    Enabled = WorkerOptions.Policy.Targeting.Enabled,
                    MaxSboms = WorkerOptions.Policy.Targeting.MaxSboms,
                    DefaultUsageOnly = WorkerOptions.Policy.Targeting.DefaultUsageOnly
                }
            }
        };
    }

    private sealed class StubPolicyRunTargetingService : IPolicyRunTargetingService
    {
        public Func<PolicyRunJob, PolicyRunTargetingResult>? OnEnsureTargets { get; set; }

        public Task<PolicyRunTargetingResult> EnsureTargetsAsync(PolicyRunJob job, CancellationToken cancellationToken)
            => Task.FromResult(OnEnsureTargets?.Invoke(job) ?? PolicyRunTargetingResult.Unchanged(job));
    }

    private sealed class RecordingPolicyRunJobRepository : IPolicyRunJobRepository
    {
        public bool ReplaceCalled { get; private set; }
        public string? ExpectedLeaseOwner { get; private set; }
        public PolicyRunJob? LastJob { get; private set; }

        public Task<PolicyRunJob?> GetAsync(string tenantId, string jobId, IClientSessionHandle? session = null, CancellationToken cancellationToken = default)
            => Task.FromResult<PolicyRunJob?>(null);

        public Task<PolicyRunJob?> GetByRunIdAsync(string tenantId, string runId, IClientSessionHandle? session = null, CancellationToken cancellationToken = default)
            => Task.FromResult<PolicyRunJob?>(null);

        public Task InsertAsync(PolicyRunJob job, IClientSessionHandle? session = null, CancellationToken cancellationToken = default)
        {
            LastJob = job;
            return Task.CompletedTask;
        }

        public Task<PolicyRunJob?> LeaseAsync(string leaseOwner, DateTimeOffset now, TimeSpan leaseDuration, int maxAttempts, IClientSessionHandle? session = null, CancellationToken cancellationToken = default)
            => Task.FromResult<PolicyRunJob?>(null);

        public Task<bool> ReplaceAsync(PolicyRunJob job, string? expectedLeaseOwner = null, IClientSessionHandle? session = null, CancellationToken cancellationToken = default)
        {
            ReplaceCalled = true;
            ExpectedLeaseOwner = expectedLeaseOwner;
            LastJob = job;
            return Task.FromResult(true);
        }

        public Task<IReadOnlyList<PolicyRunJob>> ListAsync(string tenantId, string? policyId = null, PolicyRunMode? mode = null, IReadOnlyCollection<PolicyRunJobStatus>? statuses = null, DateTimeOffset? queuedAfter = null, int limit = 50, IClientSessionHandle? session = null, CancellationToken cancellationToken = default)
            => Task.FromResult<IReadOnlyList<PolicyRunJob>>(Array.Empty<PolicyRunJob>());
    }

    private sealed class StubPolicyRunClient : IPolicyRunClient
    {
        public PolicyRunSubmissionResult Result { get; set; } = PolicyRunSubmissionResult.Succeeded(null, null);

        public Task<PolicyRunSubmissionResult> SubmitAsync(PolicyRunJob job, PolicyRunRequest request, CancellationToken cancellationToken)
            => Task.FromResult(Result);
    }

    private sealed class TestTimeProvider : TimeProvider
    {
        private DateTimeOffset _now;

        public TestTimeProvider(DateTimeOffset now)
        {
            _now = now;
        }

        public override DateTimeOffset GetUtcNow() => _now;

        public void Advance(TimeSpan delta) => _now = _now.Add(delta);
    }
}