package mongo

import (
	"context"
	"time"

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"

	"github.com/andydixon/relata/internal/domain"
)

// JobRepo is the persistent, leased job queue.
type JobRepo struct{ base }

var _ domain.JobRepository = (*JobRepo)(nil)

// maxJobHistory bounds how many attempt records one job retains.
const maxJobHistory = 10

// Enqueue adds a job to the queue.
//
// When the job carries an idempotency key, a unique partial index on pending
// jobs makes a duplicate enqueue a no-op rather than a second unit of work.
// That is what lets the pipeline re-enqueue freely after a crash.
func (r *JobRepo) Enqueue(ctx context.Context, j *domain.Job) error {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	if j.WorkspaceID == "" {
		return domain.ErrWorkspaceMiss
	}
	if !j.Type.Valid() {
		return domain.ErrInvalidInput
	}
	now := time.Now().UTC()
	if j.CreatedAt.IsZero() {
		j.CreatedAt = now
	}
	j.UpdatedAt = now
	if j.RunAfter.IsZero() {
		j.RunAfter = now
	}
	if j.Status == "" {
		j.Status = domain.JobPending
	}

	_, err := r.coll.InsertOne(ctx, j)
	if isDup(err) {
		// Equivalent work is already queued.
		r.observe("enqueue", started, nil)
		return nil
	}
	r.observe("enqueue", started, err)
	return translate(err)
}

// Claim atomically leases the next runnable job.
//
// A lease, not a lock: the document records who holds it and until when. A
// worker that dies stops renewing, the lease expires, and ReleaseExpired
// returns the job to the pool. Nothing external is needed to detect the crash.
func (r *JobRepo) Claim(ctx context.Context, owner string, types []domain.JobType, lease time.Duration) (*domain.Job, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	now := time.Now().UTC()
	f := bson.D{
		{Key: "status", Value: string(domain.JobPending)},
		{Key: "run_after", Value: bson.D{{Key: "$lte", Value: now}}},
	}
	if len(types) > 0 {
		vals := make([]any, 0, len(types))
		for _, t := range types {
			vals = append(vals, string(t))
		}
		f = append(f, bson.E{Key: "type", Value: bson.D{{Key: "$in", Value: vals}}})
	}

	expires := now.Add(lease)
	update := bson.D{
		{Key: "$set", Value: bson.D{
			{Key: "status", Value: string(domain.JobRunning)},
			{Key: "lease_owner", Value: owner},
			{Key: "lease_expires", Value: expires},
			{Key: "updated_at", Value: now},
		}},
		{Key: "$inc", Value: bson.D{{Key: "attempts", Value: 1}}},
	}

	opts := options.FindOneAndUpdate().
		SetSort(bson.D{{Key: "priority", Value: -1}, {Key: "run_after", Value: 1}}).
		SetReturnDocument(options.After)

	var j domain.Job
	err := r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&j)
	if err == mongo.ErrNoDocuments {
		r.observe("claim", started, nil)
		return nil, nil // an empty queue is not an error
	}
	r.observe("claim", started, err)
	if err != nil {
		return nil, translate(err)
	}
	return &j, nil
}

// Heartbeat extends a lease. It fails if the lease was lost — for instance
// because it expired and another worker claimed the job — which is the signal
// for the original worker to abandon its work rather than write a stale result.
func (r *JobRepo) Heartbeat(ctx context.Context, id, owner string, lease time.Duration) error {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	now := time.Now().UTC()
	f := bson.D{
		{Key: "_id", Value: id},
		{Key: "lease_owner", Value: owner},
		{Key: "status", Value: string(domain.JobRunning)},
	}
	update := bson.D{{Key: "$set", Value: bson.D{
		{Key: "lease_expires", Value: now.Add(lease)},
		{Key: "updated_at", Value: now},
	}}}

	res, err := r.coll.UpdateOne(ctx, f, update)
	r.observe("heartbeat", started, err)
	if err != nil {
		return translate(err)
	}
	if res.MatchedCount == 0 {
		return domain.ErrConflict
	}
	return nil
}

// Complete marks a job succeeded. It is scoped to the lease owner so a worker
// that lost its lease cannot mark someone else's attempt as done.
func (r *JobRepo) Complete(ctx context.Context, id, owner string, att domain.JobAttempt) error {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	now := time.Now().UTC()
	f := bson.D{{Key: "_id", Value: id}, {Key: "lease_owner", Value: owner}}
	update := bson.D{
		{Key: "$set", Value: bson.D{
			{Key: "status", Value: string(domain.JobSucceeded)},
			{Key: "completed_at", Value: now},
			{Key: "updated_at", Value: now},
			{Key: "last_error", Value: ""},
			{Key: "lease_owner", Value: ""},
		}},
		{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
		{Key: "$push", Value: bson.D{{Key: "history", Value: bson.D{
			{Key: "$each", Value: bson.A{att}},
			{Key: "$slice", Value: -maxJobHistory},
		}}}},
	}

	res, err := r.coll.UpdateOne(ctx, f, update)
	r.observe("complete", started, err)
	if err != nil {
		return translate(err)
	}
	if res.MatchedCount == 0 {
		return domain.ErrConflict
	}
	return nil
}

// Fail records a failed attempt.
//
// When retryAfter is non-nil the job returns to pending for another attempt;
// otherwise it is terminal. The decision belongs to the worker, which knows
// whether the error was transient — the queue does not guess.
func (r *JobRepo) Fail(ctx context.Context, id, owner string, att domain.JobAttempt, retryAfter *time.Time) error {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	now := time.Now().UTC()
	set := bson.D{
		{Key: "updated_at", Value: now},
		{Key: "last_error", Value: att.Error},
		{Key: "lease_owner", Value: ""},
	}
	if retryAfter != nil {
		set = append(set,
			bson.E{Key: "status", Value: string(domain.JobPending)},
			bson.E{Key: "run_after", Value: *retryAfter},
		)
	} else {
		set = append(set,
			bson.E{Key: "status", Value: string(domain.JobFailed)},
			bson.E{Key: "completed_at", Value: now},
		)
	}

	f := bson.D{{Key: "_id", Value: id}, {Key: "lease_owner", Value: owner}}
	update := bson.D{
		{Key: "$set", Value: set},
		{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
		{Key: "$push", Value: bson.D{{Key: "history", Value: bson.D{
			{Key: "$each", Value: bson.A{att}},
			{Key: "$slice", Value: -maxJobHistory},
		}}}},
	}

	res, err := r.coll.UpdateOne(ctx, f, update)
	r.observe("fail", started, err)
	if err != nil {
		return translate(err)
	}
	if res.MatchedCount == 0 {
		return domain.ErrConflict
	}
	return nil
}

// Get fetches a job by ID.
//
// Jobs are infrastructure records rather than user content, so this is not
// workspace-scoped; callers that expose a job to a user check its workspace.
func (r *JobRepo) Get(ctx context.Context, id string) (*domain.Job, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	var j domain.Job
	err := r.coll.FindOne(ctx, bson.D{{Key: "_id", Value: id}}).Decode(&j)
	r.observe("get", started, err)
	if err != nil {
		return nil, translate(err)
	}
	return &j, nil
}

// List returns a page of jobs in a workspace.
func (r *JobRepo) List(ctx context.Context, workspaceID string, statuses []domain.JobStatus, types []domain.JobType, p domain.Page) ([]*domain.Job, int64, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	f, err := wsFilter(workspaceID)
	if err != nil {
		return nil, 0, err
	}
	if len(statuses) > 0 {
		vals := make([]any, 0, len(statuses))
		for _, s := range statuses {
			vals = append(vals, string(s))
		}
		f = append(f, bson.E{Key: "status", Value: bson.D{{Key: "$in", Value: vals}}})
	}
	if len(types) > 0 {
		vals := make([]any, 0, len(types))
		for _, t := range types {
			if t.Valid() {
				vals = append(vals, string(t))
			}
		}
		if len(vals) > 0 {
			f = append(f, bson.E{Key: "type", Value: bson.D{{Key: "$in", Value: vals}}})
		}
	}
	p = p.Normalise(50, 200)

	total, err := r.coll.CountDocuments(ctx, f)
	if err != nil {
		r.observe("list", started, err)
		return nil, 0, translate(err)
	}
	cur, err := r.coll.Find(ctx, f, options.Find().
		SetSort(bson.D{{Key: "created_at", Value: -1}}).
		SetSkip(int64(p.Offset)).SetLimit(int64(p.Limit)))
	r.observe("list", started, err)
	if err != nil {
		return nil, 0, translate(err)
	}
	var out []*domain.Job
	if err := cur.All(ctx, &out); err != nil {
		return nil, 0, err
	}
	return out, total, nil
}

// Retry returns a failed job to the queue and resets its attempt budget.
func (r *JobRepo) Retry(ctx context.Context, workspaceID, id string) (*domain.Job, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	f, err := wsFilter(workspaceID)
	if err != nil {
		return nil, err
	}
	f = append(f,
		bson.E{Key: "_id", Value: id},
		bson.E{Key: "status", Value: bson.D{{Key: "$in", Value: bson.A{
			string(domain.JobFailed), string(domain.JobCancelled),
		}}}},
	)

	now := time.Now().UTC()
	update := bson.D{{Key: "$set", Value: bson.D{
		{Key: "status", Value: string(domain.JobPending)},
		{Key: "run_after", Value: now},
		{Key: "attempts", Value: 0},
		{Key: "lease_owner", Value: ""},
		{Key: "updated_at", Value: now},
	}}}

	opts := options.FindOneAndUpdate().SetReturnDocument(options.After)
	var j domain.Job
	err = r.coll.FindOneAndUpdate(ctx, f, update, opts).Decode(&j)
	r.observe("retry", started, err)
	if err != nil {
		return nil, translate(err)
	}
	return &j, nil
}

// CountByStatus groups jobs by status, for the dashboard and metrics.
func (r *JobRepo) CountByStatus(ctx context.Context, workspaceID string) (map[domain.JobStatus]int64, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	f, err := wsFilter(workspaceID)
	if err != nil {
		return nil, err
	}
	cur, err := r.coll.Aggregate(ctx, mongo.Pipeline{
		{{Key: "$match", Value: f}},
		{{Key: "$group", Value: bson.D{
			{Key: "_id", Value: "$status"},
			{Key: "n", Value: bson.D{{Key: "$sum", Value: 1}}},
		}}},
	})
	r.observe("count_by_status", started, err)
	if err != nil {
		return nil, translate(err)
	}
	defer cur.Close(ctx)

	out := map[domain.JobStatus]int64{}
	for cur.Next(ctx) {
		var row struct {
			ID string `bson:"_id"`
			N  int64  `bson:"n"`
		}
		if err := cur.Decode(&row); err == nil {
			out[domain.JobStatus(row.ID)] = row.N
		}
	}
	return out, cur.Err()
}

// ReleaseExpired returns jobs whose lease elapsed to the pending pool.
//
// This is the crash-recovery path. A job that has exhausted its attempts is
// marked failed instead, so a job that reliably kills its worker cannot loop
// forever.
func (r *JobRepo) ReleaseExpired(ctx context.Context, now time.Time) (int64, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	f := bson.D{
		{Key: "status", Value: string(domain.JobRunning)},
		{Key: "lease_expires", Value: bson.D{{Key: "$lt", Value: now}}},
	}

	// Jobs still within their attempt budget go back to pending.
	retryable := append(bson.D{}, f...)
	retryable = append(retryable, bson.E{Key: "$expr", Value: bson.D{
		{Key: "$lt", Value: bson.A{"$attempts", "$max_attempts"}},
	}})

	res, err := r.coll.UpdateMany(ctx, retryable, bson.D{
		{Key: "$set", Value: bson.D{
			{Key: "status", Value: string(domain.JobPending)},
			{Key: "run_after", Value: now},
			{Key: "lease_owner", Value: ""},
			{Key: "last_error", Value: "worker lease expired; the job was returned to the queue"},
			{Key: "updated_at", Value: now},
		}},
		{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
	})
	if err != nil {
		r.observe("release_expired", started, err)
		return 0, translate(err)
	}
	released := res.ModifiedCount

	// Jobs out of attempts are terminal.
	exhausted := append(bson.D{}, f...)
	exhausted = append(exhausted, bson.E{Key: "$expr", Value: bson.D{
		{Key: "$gte", Value: bson.A{"$attempts", "$max_attempts"}},
	}})
	_, err = r.coll.UpdateMany(ctx, exhausted, bson.D{
		{Key: "$set", Value: bson.D{
			{Key: "status", Value: string(domain.JobFailed)},
			{Key: "lease_owner", Value: ""},
			{Key: "completed_at", Value: now},
			{Key: "last_error", Value: "worker lease expired and the job has no attempts remaining"},
			{Key: "updated_at", Value: now},
		}},
		{Key: "$unset", Value: bson.D{{Key: "lease_expires", Value: ""}}},
	})
	r.observe("release_expired", started, err)
	if err != nil {
		return released, translate(err)
	}
	return released, nil
}

// DeleteForItem removes an item's queued work, used during deletion.
func (r *JobRepo) DeleteForItem(ctx context.Context, workspaceID, itemID string) (int64, error) {
	ctx, cancel := r.ctx(ctx)
	defer cancel()
	started := time.Now()

	f, err := wsFilter(workspaceID)
	if err != nil {
		return 0, err
	}
	f = append(f,
		bson.E{Key: "item_id", Value: itemID},
		// A running job is left alone; it will notice the item is gone.
		bson.E{Key: "status", Value: bson.D{{Key: "$ne", Value: string(domain.JobRunning)}}},
		// The deletion job itself must survive.
		bson.E{Key: "type", Value: bson.D{{Key: "$ne", Value: string(domain.JobDeleteItemData)}}},
	)

	res, err := r.coll.DeleteMany(ctx, f)
	r.observe("delete_for_item", started, err)
	if err != nil {
		return 0, translate(err)
	}
	return res.DeletedCount, nil
}