package automation

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/store/model"

	"gorm.io/gorm"
)

func (s *Store) EnqueueJob(ctx context.Context, job AutomationJob) (AutomationJob, error) {
	if err := s.ensureDB(); err != nil {
		return AutomationJob{}, err
	}
	job, err := enqueueAutomationJobTx(s.db.WithContext(ctx), job)
	if err != nil {
		return AutomationJob{}, err
	}
	return job, nil
}

// ScheduleJobRun marks a job run queued and creates its queue row in one
// transaction so scheduler retries never observe only one side of the state.
func (s *Store) ScheduleJobRun(ctx context.Context, jobRun AutomationJobRun, job AutomationJob) (AutomationJob, error) {
	if err := s.ensureDB(); err != nil {
		return AutomationJob{}, err
	}
	var scheduled AutomationJob
	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := tx.Save(&jobRun).Error; err != nil {
			return fmt.Errorf("update automation job run: %w", err)
		}
		job, err := enqueueAutomationJobTx(tx, job)
		if err != nil {
			return err
		}
		scheduled = job
		return nil
	})
	if err != nil {
		return AutomationJob{}, err
	}
	return scheduled, nil
}

func enqueueAutomationJobTx(tx *gorm.DB, job AutomationJob) (AutomationJob, error) {
	NormalizeAutomationJobClaimFields(&job)
	var existing AutomationJob
	err := tx.Where("idempotency_key = ?", job.IdempotencyKey).First(&existing).Error
	if err == nil {
		return existing, nil
	}
	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return AutomationJob{}, fmt.Errorf("get existing automation job: %w", err)
	}
	if err := tx.Create(&job).Error; err != nil {
		return AutomationJob{}, fmt.Errorf("enqueue automation job: %w", err)
	}
	return job, nil
}

type ClaimJobsRequest struct {
	WorkspaceID     string
	QueueName       string
	ShardBucket     int
	ShardBuckets    []int
	LeaseOwner      string
	LeaseDuration   time.Duration
	Limit           int
	RunnerGroup     string
	Labels          []string
	MaxRunnerLabels int
	Now             time.Time
}

type AutomationJobFailureResult struct {
	JobStatus       string
	RetryScheduled  bool
	AttemptsExhaust bool
}

func RetryBackoff(attempt int) time.Duration {
	if attempt < 1 {
		attempt = 1
	}
	backoff := 30 * time.Second
	for i := 1; i < attempt; i++ {
		backoff *= 2
		if backoff >= 5*time.Minute {
			return 5 * time.Minute
		}
	}
	return backoff
}

func (s *Store) GetJob(ctx context.Context, jobID int64) (AutomationJob, error) {
	if err := s.ensureDB(); err != nil {
		return AutomationJob{}, err
	}
	var job AutomationJob
	err := s.db.WithContext(ctx).Where("id = ?", jobID).First(&job).Error
	if err != nil {
		return AutomationJob{}, err
	}
	return job, nil
}

func (s *Store) ClaimJobs(ctx context.Context, request ClaimJobsRequest) ([]AutomationJob, error) {
	if err := s.ensureDB(); err != nil {
		return nil, err
	}
	if request.Limit <= 0 {
		request.Limit = 1
	}
	if request.LeaseDuration <= 0 {
		request.LeaseDuration = time.Minute
	}
	workspaceID := strings.TrimSpace(request.WorkspaceID)
	if workspaceID == "" {
		return nil, fmt.Errorf("automation claim workspace id is required")
	}
	maxRunnerLabels := request.MaxRunnerLabels
	if maxRunnerLabels <= 0 || maxRunnerLabels > automationClaimMaxRunnerLabels {
		maxRunnerLabels = automationClaimMaxRunnerLabels
	}
	runnerLabels := normalizeAutomationClaimLabels(request.Labels)
	if len(runnerLabels) > maxRunnerLabels {
		return nil, fmt.Errorf("automation runner labels exceed maximum count")
	}
	shardBuckets, err := automationClaimShardBuckets(request)
	if err != nil {
		return nil, err
	}
	runnerGroupKeys := []string{""}
	if runnerGroupKey := normalizeAutomationClaimToken(request.RunnerGroup); runnerGroupKey != "" {
		runnerGroupKeys = append(runnerGroupKeys, runnerGroupKey)
	}
	labelKeys := automationClaimSubsetKeys(runnerLabels)
	labelKeyHashes := automationClaimKeyHashes(labelKeys)
	leaseUntil := request.Now.Add(request.LeaseDuration)
	var claimed []AutomationJob
	err = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var candidates []AutomationJob
		if err := tx.Where("workspace_id = ? AND queue_name = ? AND shard_bucket IN ? AND status = ? AND runner_group_key IN ? AND label_key_hash IN ? AND label_key IN ? AND scheduled_at <= ?", workspaceID, request.QueueName, shardBuckets, AutomationJobStatusReady, runnerGroupKeys, labelKeyHashes, labelKeys, request.Now).Order("scheduled_at ASC").Order("priority ASC").Order("id ASC").Limit(request.Limit).Find(&candidates).Error; err != nil {
			return err
		}
		for _, candidate := range candidates {
			if len(claimed) >= request.Limit {
				break
			}
			result := tx.Model(&AutomationJob{}).Where("id = ? AND workspace_id = ? AND status = ?", candidate.ID, workspaceID, AutomationJobStatusReady).Updates(map[string]any{
				"status":      AutomationJobStatusLeased,
				"lease_owner": request.LeaseOwner,
				"lease_until": leaseUntil,
				"attempt":     candidate.Attempt + 1,
				"updated_at":  request.Now,
			})
			if result.Error != nil {
				return result.Error
			}
			if result.RowsAffected == 0 {
				continue
			}
			candidate.Status = AutomationJobStatusLeased
			candidate.LeaseOwner = request.LeaseOwner
			candidate.LeaseUntil = &leaseUntil
			candidate.Attempt++
			claimed = append(claimed, candidate)
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("claim automation jobs: %w", err)
	}
	return claimed, nil
}

func automationClaimShardBuckets(request ClaimJobsRequest) ([]int, error) {
	raw := request.ShardBuckets
	if len(raw) == 0 {
		raw = []int{request.ShardBucket}
	}
	out := make([]int, 0, len(raw))
	seen := map[int]struct{}{}
	for _, shard := range raw {
		if shard < 0 || shard >= AutomationClaimShardCount {
			return nil, fmt.Errorf("automation claim shard_bucket is invalid")
		}
		if _, ok := seen[shard]; ok {
			continue
		}
		seen[shard] = struct{}{}
		out = append(out, shard)
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("automation claim shard_bucket is required")
	}
	return out, nil
}

// NormalizeAutomationJobClaimFields fills denormalized queue-claim columns from
// runner group and labels fields.
func NormalizeAutomationJobClaimFields(job *AutomationJob) {
	if job == nil {
		return
	}
	job.RunnerGroupKey = normalizeAutomationClaimToken(job.RunnerGroup)
	labels := normalizeAutomationClaimLabels(decodeLabelsJSON(job.LabelsJSON))
	job.LabelKey = strings.Join(labels, ",")
	job.LabelKeyHash = automationClaimKeyHash(job.LabelKey)
	job.LabelCount = len(labels)
}

func normalizeAutomationClaimLabels(labels []string) []string {
	if len(labels) == 0 {
		return nil
	}
	out := make([]string, 0, len(labels))
	seen := map[string]struct{}{}
	for _, label := range labels {
		normalized := normalizeAutomationClaimToken(label)
		if normalized == "" {
			continue
		}
		if _, ok := seen[normalized]; ok {
			continue
		}
		seen[normalized] = struct{}{}
		out = append(out, normalized)
	}
	sort.Strings(out)
	return out
}

func automationClaimSubsetKeys(labels []string) []string {
	if len(labels) == 0 {
		return []string{""}
	}
	subsetCount := 1 << len(labels)
	keys := make([]string, 0, subsetCount)
	for mask := 0; mask < subsetCount; mask++ {
		subset := make([]string, 0, len(labels))
		for idx, label := range labels {
			if mask&(1<<idx) != 0 {
				subset = append(subset, label)
			}
		}
		keys = append(keys, strings.Join(subset, ","))
	}
	return keys
}

func automationClaimKeyHashes(labelKeys []string) []string {
	hashes := make([]string, 0, len(labelKeys))
	for _, key := range labelKeys {
		hashes = append(hashes, automationClaimKeyHash(key))
	}
	return hashes
}

func automationClaimKeyHash(key string) string {
	key = strings.TrimSpace(key)
	if key == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(key))
	return hex.EncodeToString(sum[:])
}

func normalizeAutomationClaimToken(value string) string {
	return strings.ToLower(strings.TrimSpace(value))
}

func decodeLabelsJSON(raw string) []string {
	var labels []string
	if strings.TrimSpace(raw) == "" {
		return nil
	}
	if err := model.UnmarshalJSONString("automation runner labels", raw, &labels); err != nil {
		return nil
	}
	out := labels[:0]
	for _, label := range labels {
		if trimmed := strings.TrimSpace(label); trimmed != "" {
			out = append(out, trimmed)
		}
	}
	return out
}

func (s *Store) MarkJobAssignmentRunning(ctx context.Context, jobID int64, leaseOwner string, jobRun AutomationJobRun, now time.Time) error {
	if err := s.ensureDB(); err != nil {
		return err
	}
	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if err := verifyLeasedAutomationJobTx(tx, jobID, leaseOwner, now); err != nil {
			return err
		}
		if err := tx.Save(&jobRun).Error; err != nil {
			return fmt.Errorf("mark automation job running: %w", err)
		}
		return nil
	})
}

func (s *Store) CompleteJobAssignmentSuccess(ctx context.Context, jobID int64, leaseOwner string, jobRun AutomationJobRun, event AutomationEvent, now time.Time) error {
	if err := s.ensureDB(); err != nil {
		return err
	}
	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		result := tx.Model(&AutomationJob{}).Where("id = ? AND status = ? AND lease_owner = ? AND lease_until >= ?", jobID, AutomationJobStatusLeased, strings.TrimSpace(leaseOwner), now).Updates(map[string]any{
			"status":      AutomationJobStatusDone,
			"lease_owner": "",
			"lease_until": nil,
			"updated_at":  now,
		})
		if result.Error != nil {
			return fmt.Errorf("complete automation job: %w", result.Error)
		}
		if result.RowsAffected == 0 {
			return gorm.ErrRecordNotFound
		}
		if err := tx.Save(&jobRun).Error; err != nil {
			return fmt.Errorf("update automation job success: %w", err)
		}
		if err := tx.Create(&event).Error; err != nil {
			return fmt.Errorf("create automation success event: %w", err)
		}
		return nil
	})
}

func (s *Store) FailJobAssignment(ctx context.Context, jobID int64, leaseOwner string, jobRun AutomationJobRun, event AutomationEvent, scheduledAt time.Time, now time.Time, terminal bool) (AutomationJobFailureResult, error) {
	if err := s.ensureDB(); err != nil {
		return AutomationJobFailureResult{}, err
	}
	var result AutomationJobFailureResult
	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var job AutomationJob
		if err := tx.Where("id = ? AND status = ? AND lease_owner = ? AND lease_until >= ?", jobID, AutomationJobStatusLeased, strings.TrimSpace(leaseOwner), now).First(&job).Error; err != nil {
			return err
		}
		status := AutomationJobStatusReady
		result.RetryScheduled = true
		if terminal || job.Attempt >= job.MaxAttempts {
			status = AutomationJobStatusFailed
			result.RetryScheduled = false
			result.AttemptsExhaust = true
		}
		update := tx.Model(&AutomationJob{}).Where("id = ? AND status = ? AND lease_owner = ?", jobID, AutomationJobStatusLeased, strings.TrimSpace(leaseOwner)).Updates(map[string]any{
			"status":       status,
			"lease_owner":  "",
			"lease_until":  nil,
			"scheduled_at": scheduledAt,
			"updated_at":   now,
		})
		if update.Error != nil {
			return fmt.Errorf("fail automation job: %w", update.Error)
		}
		if update.RowsAffected == 0 {
			return gorm.ErrRecordNotFound
		}
		if err := tx.Save(&jobRun).Error; err != nil {
			return fmt.Errorf("update automation job failure: %w", err)
		}
		if err := tx.Create(&event).Error; err != nil {
			return fmt.Errorf("create automation failure event: %w", err)
		}
		result.JobStatus = status
		return nil
	})
	if err != nil {
		return AutomationJobFailureResult{}, err
	}
	return result, nil
}