package workerruntime

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// ErrLeaseLost reports that the job run row this worker owned was completed or
// reclaimed elsewhere, so this worker's result must not overwrite it.
var ErrLeaseLost = errors.New("worker job lease was lost")

// Lease identifies one claimed scheduling window.
type Lease struct {
	ID          string
	JobType     string
	JobKey      string
	ScheduledAt time.Time
	Owner       string
}

// LeaseStore coordinates duplicate workers through the database.
//
// Acquire must be atomic per (jobType, jobKey, scheduledAt): exactly one caller
// may hold a window at a time, an already completed window must never be handed
// out again, and a window whose lease has expired must be reclaimable so a
// crashed worker cannot block a job forever.
type LeaseStore interface {
	Acquire(ctx context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error)
	Complete(ctx context.Context, lease Lease, status, errorCode string, counts map[string]int64) error
}

// PostgresLeaseStore implements LeaseStore on the existing job_runs table.
//
// It follows the same pattern as alertworker.PostgresLeaseStore, which is
// hardwired to the alert evaluation job type; this one is parameterized by job
// type and additionally records result counts, so every worker job shares one
// coordination and reporting mechanism instead of inventing its own.
type PostgresLeaseStore struct {
	Pool *pgxpool.Pool
}

// Acquire claims one scheduling window, reclaiming an expired lease if needed.
func (s PostgresLeaseStore) Acquire(ctx context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error) {
	if s.Pool == nil {
		return Lease{}, false, fmt.Errorf("%w: lease store has no database pool", ErrInvalidConfig)
	}
	if strings.TrimSpace(jobType) == "" || strings.TrimSpace(jobKey) == "" || strings.TrimSpace(owner) == "" || ttl <= 0 {
		return Lease{}, false, fmt.Errorf("%w: lease identity is invalid", ErrInvalidConfig)
	}
	scheduled := scheduledAt.UTC()
	moment := now.UTC()
	leaseUntil := moment.Add(ttl)
	tx, err := s.Pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return Lease{}, false, fmt.Errorf("begin worker lease: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	leaseID := newID()
	var insertedID string
	err = tx.QueryRow(ctx, `INSERT INTO job_runs (id,job_type,job_key,scheduled_at,started_at,status,lease_owner,lease_until) VALUES ($1,$2,$3,$4,$5,'running',$6,$7) ON CONFLICT (job_type,job_key,scheduled_at) DO NOTHING RETURNING id`,
		leaseID, jobType, jobKey, scheduled, moment, owner, leaseUntil).Scan(&insertedID)
	if err == nil {
		if err := tx.Commit(ctx); err != nil {
			return Lease{}, false, fmt.Errorf("commit worker lease: %w", err)
		}
		return Lease{ID: insertedID, JobType: jobType, JobKey: jobKey, ScheduledAt: scheduled, Owner: owner}, true, nil
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return Lease{}, false, fmt.Errorf("insert worker lease: %w", err)
	}
	var status string
	var existingUntil *time.Time
	if err := tx.QueryRow(ctx, `SELECT status,lease_until FROM job_runs WHERE job_type=$1 AND job_key=$2 AND scheduled_at=$3 FOR UPDATE`, jobType, jobKey, scheduled).Scan(&status, &existingUntil); err != nil {
		return Lease{}, false, fmt.Errorf("read worker lease: %w", err)
	}
	if status != "running" || (existingUntil != nil && existingUntil.After(moment)) {
		if err := tx.Commit(ctx); err != nil {
			return Lease{}, false, fmt.Errorf("commit worker lease contention: %w", err)
		}
		return Lease{}, false, nil
	}
	tag, err := tx.Exec(ctx, `UPDATE job_runs SET started_at=$1,completed_at=NULL,error_code=NULL,lease_owner=$2,lease_until=$3 WHERE job_type=$4 AND job_key=$5 AND scheduled_at=$6 AND status='running' AND (lease_until IS NULL OR lease_until <= $7)`,
		moment, owner, leaseUntil, jobType, jobKey, scheduled, moment)
	if err != nil {
		return Lease{}, false, fmt.Errorf("reclaim worker lease: %w", err)
	}
	if tag.RowsAffected() != 1 {
		if err := tx.Commit(ctx); err != nil {
			return Lease{}, false, fmt.Errorf("commit worker lease reclaim contention: %w", err)
		}
		return Lease{}, false, nil
	}
	if err := tx.Commit(ctx); err != nil {
		return Lease{}, false, fmt.Errorf("commit worker lease reclaim: %w", err)
	}
	return Lease{ID: leaseID, JobType: jobType, JobKey: jobKey, ScheduledAt: scheduled, Owner: owner}, true, nil
}

// Complete records the outcome of a claimed window. The update is scoped to the
// owning worker, so a run whose lease already expired cannot overwrite the
// result of the worker that took the window over.
func (s PostgresLeaseStore) Complete(ctx context.Context, lease Lease, status, errorCode string, counts map[string]int64) error {
	if s.Pool == nil {
		return fmt.Errorf("%w: lease store has no database pool", ErrInvalidConfig)
	}
	if !recordedStatus(status) {
		return fmt.Errorf("%w: worker job status %q is invalid", ErrInvalidConfig, status)
	}
	payload, err := encodeCounts(status, errorCode, counts)
	if err != nil {
		return err
	}
	// error_code describes a failure; the reason for a skipped or disabled run
	// lives in counts so a non-failure is never read back as an error.
	failureCode := ""
	if status == StatusFailed {
		failureCode = boundedCode(errorCode)
	}
	tag, err := s.Pool.Exec(ctx, `UPDATE job_runs SET status=$1,completed_at=now(),counts=$2::jsonb,error_code=NULLIF($3,''),lease_owner=NULL,lease_until=NULL WHERE job_type=$4 AND job_key=$5 AND scheduled_at=$6 AND lease_owner=$7 AND status='running'`,
		status, payload, failureCode, lease.JobType, lease.JobKey, lease.ScheduledAt.UTC(), lease.Owner)
	if err != nil {
		return fmt.Errorf("complete worker job: %w", err)
	}
	if tag.RowsAffected() != 1 {
		return ErrLeaseLost
	}
	return nil
}

// recordedStatus bounds what may be written to job_runs.status. Skipped and
// disabled are kept distinct from completed on purpose: a run that did nothing
// must never be read back as a successful observation.
func recordedStatus(status string) bool {
	switch status {
	case StatusCompleted, StatusFailed, StatusSkipped, StatusDisabled:
		return true
	default:
		return false
	}
}

func encodeCounts(status, reason string, counts map[string]int64) ([]byte, error) {
	document := make(map[string]any, len(counts)+2)
	for key, value := range counts {
		if len(document) >= 20 {
			break
		}
		document[boundedCode(key)] = value
	}
	document["status"] = status
	if reason != "" {
		document["reason"] = boundedCode(reason)
	}
	payload, err := json.Marshal(document)
	if err != nil {
		return nil, fmt.Errorf("encode worker job counts: %w", err)
	}
	return payload, nil
}

func boundedCode(value string) string {
	value = strings.TrimSpace(value)
	if len(value) > 160 {
		return value[:160]
	}
	return value
}

// MemoryLeaseStore is an in-memory LeaseStore with the same semantics, used by
// tests and by a worker running without a database.
type MemoryLeaseStore struct {
	mu      sync.Mutex
	records map[string]memoryLease
}

type memoryLease struct {
	lease      Lease
	status     string
	leaseUntil time.Time
	counts     map[string]int64
	reason     string
}

// NewMemoryLeaseStore returns an empty in-memory lease store.
func NewMemoryLeaseStore() *MemoryLeaseStore {
	return &MemoryLeaseStore{records: make(map[string]memoryLease)}
}

// Acquire claims a window unless it is completed or still leased.
func (s *MemoryLeaseStore) Acquire(_ context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error) {
	if strings.TrimSpace(jobType) == "" || strings.TrimSpace(jobKey) == "" || strings.TrimSpace(owner) == "" || ttl <= 0 {
		return Lease{}, false, fmt.Errorf("%w: lease identity is invalid", ErrInvalidConfig)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.records == nil {
		s.records = make(map[string]memoryLease)
	}
	key := memoryLeaseKey(jobType, jobKey, scheduledAt)
	if record, exists := s.records[key]; exists {
		if record.status != StatusRunning || record.leaseUntil.After(now.UTC()) {
			return Lease{}, false, nil
		}
	}
	lease := Lease{ID: newID(), JobType: jobType, JobKey: jobKey, ScheduledAt: scheduledAt.UTC(), Owner: owner}
	s.records[key] = memoryLease{lease: lease, status: StatusRunning, leaseUntil: now.UTC().Add(ttl)}
	return lease, true, nil
}

// Complete records the outcome for a window this owner still holds.
func (s *MemoryLeaseStore) Complete(_ context.Context, lease Lease, status, errorCode string, counts map[string]int64) error {
	if !recordedStatus(status) {
		return fmt.Errorf("%w: worker job status %q is invalid", ErrInvalidConfig, status)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	key := memoryLeaseKey(lease.JobType, lease.JobKey, lease.ScheduledAt)
	record, exists := s.records[key]
	if !exists || record.lease.ID != lease.ID || record.lease.Owner != lease.Owner || record.status != StatusRunning {
		return ErrLeaseLost
	}
	record.status = status
	record.reason = errorCode
	record.counts = counts
	record.leaseUntil = time.Time{}
	s.records[key] = record
	return nil
}

// Status reports the recorded status of one window, or an empty string when the
// window was never claimed.
func (s *MemoryLeaseStore) Status(jobType, jobKey string, scheduledAt time.Time) string {
	s.mu.Lock()
	defer s.mu.Unlock()
	record, exists := s.records[memoryLeaseKey(jobType, jobKey, scheduledAt)]
	if !exists {
		return ""
	}
	return record.status
}

// Windows reports how many distinct scheduling windows were claimed.
func (s *MemoryLeaseStore) Windows() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.records)
}

func memoryLeaseKey(jobType, jobKey string, scheduledAt time.Time) string {
	return jobType + "|" + jobKey + "|" + scheduledAt.UTC().Format(time.RFC3339Nano)
}

func newID() string {
	raw := make([]byte, 16)
	if _, err := rand.Read(raw); err != nil {
		return fmt.Sprintf("00000000-0000-4000-8000-%012d", time.Now().UnixNano()%1_000_000_000_000)
	}
	raw[6] = (raw[6] & 0x0f) | 0x40
	raw[8] = (raw[8] & 0x3f) | 0x80
	encoded := hex.EncodeToString(raw)
	return encoded[:8] + "-" + encoded[8:12] + "-" + encoded[12:16] + "-" + encoded[16:20] + "-" + encoded[20:]
}