package runtime

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/store/automation"
)

type Scheduler struct {
	store *automation.Store
	now   func() time.Time
}

func NewScheduler(store *automation.Store, now func() time.Time) *Scheduler {
	return &Scheduler{store: store, now: now}
}

func (s *Scheduler) ScheduleRun(ctx context.Context, workspaceID string, runID string) ([]automation.AutomationJob, error) {
	if s == nil || s.store == nil {
		return nil, fmt.Errorf("automation scheduler store is nil")
	}
	now := time.Now().UTC()
	if s.now != nil {
		now = s.now()
	}
	jobRuns, err := s.store.ListJobRuns(ctx, workspaceID, runID)
	if err != nil {
		return nil, err
	}
	byID := map[string]automation.AutomationJobRun{}
	for _, jobRun := range jobRuns {
		byID[jobRun.ID] = jobRun
	}
	var jobs []automation.AutomationJob
	for _, jobRun := range jobRuns {
		waitKind, requiresServerManagedRunner := jobRuntimeTraitsFromInput(jobRun.InputJSON)
		if jobRun.Status == automation.AutomationJobRunStatusQueued {
			if waitKind != "" {
				continue
			}
			job, err := s.store.EnqueueJob(ctx, assignmentForJobRun(workspaceID, runID, jobRun, requiresServerManagedRunner, now))
			if err != nil {
				return nil, err
			}
			jobs = append(jobs, job)
			continue
		}
		if jobRun.Status != automation.AutomationJobRunStatusPending && jobRun.Status != automation.AutomationJobRunStatusReady {
			continue
		}
		var deps []string
		if jobRun.DependsOnJSON != "" {
			_ = json.Unmarshal([]byte(jobRun.DependsOnJSON), &deps)
		}
		ready := true
		for _, depID := range deps {
			if byID[depID].Status != automation.AutomationJobRunStatusSucceeded {
				ready = false
				break
			}
		}
		if !ready {
			continue
		}
		jobRun.Status = automation.AutomationJobRunStatusQueued
		jobRun.ScheduledAt = now
		jobRun.UpdatedAt = now
		if waitKind != "" {
			run, err := s.store.GetRun(ctx, workspaceID, runID)
			if err != nil {
				return nil, err
			}
			jobRun.Status = automation.AutomationJobRunStatusWaiting
			wait, event, err := initialWaitForJobRun(run, jobRun, now, randomRuntimeID)
			if err != nil {
				return nil, err
			}
			if _, err := s.store.CreateWait(ctx, wait, jobRun, event); err != nil {
				return nil, err
			}
			continue
		}
		job, err := s.store.ScheduleJobRun(ctx, jobRun, assignmentForJobRun(workspaceID, runID, jobRun, requiresServerManagedRunner, now))
		if err != nil {
			return nil, err
		}
		jobs = append(jobs, job)
	}
	return jobs, nil
}