package runner

import (
	"context"
	"errors"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/jackc/pgx/v5"
	"github.com/riverqueue/river"
	"github.com/riverqueue/river/rivertype"
)

// riverExecutionBroker is the River adapter for the platform-neutral broker
// port. Queue names and River uniqueness semantics do not escape this adapter.
type riverExecutionBroker struct {
	service *Service
}

func (broker riverExecutionBroker) Publish(
	ctx context.Context,
	poolID string,
	taskID string,
	notBefore time.Time,
) (int64, error) {
	tx, err := broker.service.store.Pool().Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer func() {
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = tx.Rollback(rollbackCtx)
	}()
	jobID, err := broker.publishTx(ctx, tx, poolID, taskID, notBefore, 2)
	if err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return jobID, nil
}

func (broker riverExecutionBroker) publishTx(
	ctx context.Context,
	tx pgx.Tx,
	poolID string,
	taskID string,
	notBefore time.Time,
	priority int,
) (int64, error) {
	riverClient := broker.service.asyncControlClient()
	if riverClient == nil {
		return 0, errors.New("River execution broker is not started")
	}
	opts := riverExecutionInsertOptions(poolID, notBefore, priority)
	result, err := riverClient.InsertTx(ctx, tx, asyncTaskArgs{TaskID: taskID}, opts)
	if err != nil {
		return 0, err
	}
	if result.Job == nil {
		return 0, nil
	}
	if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2,
    updated_at = now()
WHERE id = $1::uuid`, taskID, result.Job.ID); err != nil {
		return 0, err
	}
	return result.Job.ID, nil
}

func riverExecutionInsertOptions(poolID string, notBefore time.Time, priority int) *river.InsertOpts {
	if priority < 1 {
		priority = 2
	}
	opts := &river.InsertOpts{
		MaxAttempts: 1000,
		Priority:    priority,
		Queue:       executionpool.QueueName(poolID),
		Tags:        []string{"gateway-task"},
		UniqueOpts: river.UniqueOpts{
			ByArgs: true, ByQueue: true,
			ByState: []rivertype.JobState{
				rivertype.JobStateAvailable, rivertype.JobStatePending,
				rivertype.JobStateRetryable, rivertype.JobStateRunning,
				rivertype.JobStateScheduled,
			},
		},
	}
	if !notBefore.IsZero() {
		opts.ScheduledAt = notBefore
	}
	return opts
}

var _ executionpool.ExecutionBroker = riverExecutionBroker{}