public abstract inherited sharing class RecoverableQueueable implements Queueable {
  // Control flag to determine if the next job should be enqueued on success.
  // Defaults to true, but RetryService can set this to false.
  public Boolean chainOnSuccess = true;

  // Test-visible list to track chained jobs in unit tests without actually enqueueing them
  @TestVisible
  public static List<Queueable> testChainedJobs = new List<Queueable>();

  public void execute(QueueableContext context) {
    // Attach the finalizer immediately to capture the state for retry.
    System.attachFinalizer(new JobFinalizer(this));

    // Execute the actual business logic
    this.work(context);

    // Handle Auto-Chaining
    if (this.chainOnSuccess) {
      Queueable next = this.getNextJob();
      if (next != null) {
        enqueueChain(next);
      }
    }
  }

  // Centralized method to handle chaining, allowing interception during tests
  public static void enqueueChain(Queueable job) {
    if (Test.isRunningTest()) {
      testChainedJobs.add(job);
    } else {
      System.enqueueJob(job);
    }
  }

  // Abstract method for subclasses to implement their logic
  protected abstract void work(QueueableContext context);

  // Get class name
  public abstract String getJobName();

  // Virtual method to provide the next job in the chain.
  // Override this if you want to implement a chain that continues even if this job fails.
  public virtual Queueable getNextJob() {
    return null;
  }
}

public without sharing class JobFinalizer implements Finalizer {
  private RecoverableQueueable job;
  private String jobClassName;

  public JobFinalizer(RecoverableQueueable job) {
    this.job = job;
    // Capture the class name for deserialization later.
    this.jobClassName = job.getJobName();
  }

  public void execute(FinalizerContext context) {
    // We only care about failures (UNHANDLED_EXCEPTION)
    if (context.getResult() == ParentJobResult.UNHANDLED_EXCEPTION) {
      handleFailure(context);
    }
  }

  private void handleFailure(FinalizerContext context) {
    Exception ex = context.getException();

    Async_Job_Failure__c failure = new Async_Job_Failure__c();
    failure.Apex_Class_Name__c = this.jobClassName;
    failure.Job_ID__c = context.getAsyncApexJobId();
    failure.Error_Message__c = ex != null ? ex.getMessage() : 'Unknown Error';
    failure.Stack_Trace__c = ex != null ? ex.getStackTraceString() : '';
    failure.Payload__c = JSON.serialize(this.job);

    // Ensure field lengths are respected (assuming standard text area limits, adjust as needed)
    if (
      failure.Error_Message__c != null &&
      failure.Error_Message__c.length() > 32768
    ) {
      failure.Error_Message__c = failure.Error_Message__c.substring(0, 32768);
    }

    if (
      failure.Stack_Trace__c != null &&
      failure.Stack_Trace__c.length() > 32768
    ) {
      failure.Stack_Trace__c = failure.Stack_Trace__c.substring(0, 32768);
    }

    failure.Status__c = this.attemptChainRecovery();

    insert failure;
  }

  private String attemptChainRecovery() {
    try {
      // Even if this job failed, we attempt to enqueue the next job in the chain
      // so that the rest of the process is not blocked.
      Queueable next = this.job.getNextJob();
      if (next != null) {
        RecoverableQueueable.enqueueChain(next);
        return 'Failed - chain recovered';
      } else {
        return 'Failed - no chain to recover';
      }
    } catch (Exception e) {
      System.debug(
        LoggingLevel.ERROR,
        'Failed to recover chain: ' + e.getMessage()
      );
      return 'Failed - chain recovery failed - ' + e.getMessage();
    }
  }
}

public with sharing class RetryService {
  public static Id retryJob(Id failureId) {
    return retryJob(failureId, false); // Default: Do not continue chain on retry (assume chain continued via Finalizer)
  }

  public static Id retryJob(Id failureId, Boolean continueChain) {
    // Retrieve the failure record
    Async_Job_Failure__c failure = [
      SELECT Payload__c, Apex_Class_Name__c, Status__c
      FROM Async_Job_Failure__c
      WHERE Id = :failureId
      LIMIT 1
    ];

    // Prevent double-retries
    if (failure.Status__c == 'Retried') {
      throw new RetryException('This job has already been retried.');
    }

    // Resolve the Type
    Type jobType = Type.forName(failure.Apex_Class_Name__c);
    if (jobType == null) {
      throw new RetryException(
        'Could not find type: ' + failure.Apex_Class_Name__c
      );
    }

    // Deserialize the job state
    Object jobInstance = JSON.deserialize(failure.Payload__c, jobType);

    Id newJobId;

    // Re-enqueue the job
    if (jobInstance instanceof RecoverableQueueable) {
      RecoverableQueueable recoverable = (RecoverableQueueable) jobInstance;
      recoverable.chainOnSuccess = continueChain;
      newJobId = System.enqueueJob(recoverable);
    } else if (jobInstance instanceof Queueable) {
      // Fallback for non-recoverable queueables (shouldn't happen with this framework)
      newJobId = System.enqueueJob((Queueable) jobInstance);
    } else {
      throw new RetryException('Restored object is not a Queueable');
    }

    // Update status to prevent future retries
    failure.Status__c = 'Retried';
    failure.Retry_Job_ID__c = newJobId;
    update failure;

    return newJobId;
  }

  public class RetryException extends Exception {
  }
}