//go:generate stringer -type=JobStateType --trimprefix=JobStateType --output job_string.go
package models

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"golang.org/x/exp/maps"

	"github.com/bacalhau-project/bacalhau/pkg/lib/validate"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
)

type JobStateType int

const (
	JobStateTypeUndefined JobStateType = iota

	// JobStateTypePending is the state of a job that has been submitted but not
	// yet scheduled.
	JobStateTypePending

	// JobStateTypeQueued is the state of a job that has been evaluated but no
	// matching nodes are available yet.
	JobStateTypeQueued

	// JobStateTypeRunning is the state of a job that has been scheduled, with at
	// least one active execution.
	JobStateTypeRunning

	// JobStateTypeCompleted is the state of a job that has successfully completed.
	// Only valid for batch jobs.
	JobStateTypeCompleted

	// JobStateTypeFailed is the state of a job that has failed.
	JobStateTypeFailed

	// JobStateTypeStopped is the state of a job that has been stopped by the user.
	JobStateTypeStopped
)

// IsUndefined returns true if the job state is undefined
func (s JobStateType) IsUndefined() bool {
	return s == JobStateTypeUndefined
}

// IsTerminal returns true if the job state is terminal
func (s JobStateType) IsTerminal() bool {
	switch s {
	case JobStateTypeCompleted, JobStateTypeFailed, JobStateTypeStopped:
		return true
	default:
		return false
	}
}

// IsRerunnable returns true if the job in a state to be re-run
func (s JobStateType) IsRerunnable() bool {
	switch s {
	case JobStateTypePending, JobStateTypeQueued, JobStateTypeUndefined:
		return false
	default:
		return true
	}
}

func JobStateTypes() []JobStateType {
	var res []JobStateType
	for typ := JobStateTypePending; typ <= JobStateTypeStopped; typ++ {
		res = append(res, typ)
	}
	return res
}

func (s JobStateType) MarshalText() ([]byte, error) {
	return []byte(s.String()), nil
}

func (s *JobStateType) UnmarshalText(text []byte) (err error) {
	name := strings.TrimSpace(string(text))
	for _, typ := range JobStateTypes() {
		if strings.EqualFold(typ.String(), name) {
			*s = typ
			return
		}
	}
	return
}

type Job struct {
	// ID is a unique identifier assigned to this job.
	// It helps to distinguish jobs with the same name after they have been deleted and re-created.
	// The ID is generated by the server and should not be set directly by the client.
	ID string `json:"ID"`

	// Name is the logical name of the job used to refer to it.
	// Submitting a job with the same name as an existing job will result in an
	// update to the existing job.
	Name string `json:"Name"`

	// Namespace is the namespace this job is running in.
	Namespace string `json:"Namespace"`

	// Type is the type of job this is, e.g. "daemon" or "batch".
	Type string `json:"Type"`

	// Priority defines the scheduling priority of this job.
	Priority int `json:"Priority"`

	// Count is the number of replicas that should be scheduled.
	// For batch and service jobs:
	// - If not present in JSON, defaults to 1
	// - If explicitly set to 0, means stop all executions
	// - If > 0, specifies exact number of replicas
	// For daemon and ops jobs:
	// - Values of 0 or 1 are ignored (job runs on all matching nodes)
	// - Values > 1 are invalid and will cause validation to fail
	Count int `json:"Count"`

	// Constraints is a selector which must be true for the compute node to run this job.
	Constraints []*LabelSelectorRequirement `json:"Constraints"`

	// Meta is used to associate arbitrary metadata with this job.
	Meta map[string]string `json:"Meta"`

	// Labels is used to associate arbitrary labels with this job, which can be used
	// for filtering.
	// key=value
	Labels map[string]string `json:"Labels"`

	Tasks []*Task `json:"Tasks"`

	// State is the current state of the job.
	State State[JobStateType] `json:"State"`

	// Version is a per-job monotonically increasing version number that is incremented
	// on each job specification update.
	Version uint64 `json:"Version"`

	// Revision is a per-job monotonically increasing revision number that is incremented
	// on each update to the job's state or specification
	Revision uint64 `json:"Revision"`

	CreateTime int64 `json:"CreateTime"`
	ModifyTime int64 `json:"ModifyTime"`
}

func (j *Job) MetricAttributes() []attribute.KeyValue {
	// TODO(forrest): will need to re-think how we tag metrics from jobs with more than one task when ever that happens.
	return append(j.Task().MetricAttributes(), attribute.String("job_type", j.Type))
}

func (j *Job) String() string {
	return j.ID
}

// NamespacedID returns the namespaced id useful for logging
func (j *Job) NamespacedID() NamespacedID {
	return NamespacedID{
		ID:        j.ID,
		Namespace: j.Namespace,
	}
}

// Normalize is used to canonicalize fields in the Job. This should be
// called when registering a Job.
func (j *Job) Normalize() {
	if j == nil {
		return
	}

	// Ensure that an empty and nil map are treated the same to avoid scheduling
	// problems since we use reflect DeepEquals.
	if j.Meta == nil {
		j.Meta = make(map[string]string)
	}

	if j.Labels == nil {
		j.Labels = make(map[string]string)
	}

	if j.Constraints == nil {
		j.Constraints = make([]*LabelSelectorRequirement, 0)
	}

	if j.Tasks == nil {
		j.Tasks = make([]*Task, 0)
	}

	// Ensure the job is in a namespace.
	if j.Namespace == "" {
		j.Namespace = DefaultNamespace
	}

	if j.Name == "" {
		j.Name = j.ID
	}

	// Handle count based on job type
	switch j.Type {
	case JobTypeDaemon, JobTypeOps:
		j.Count = 0 // Always 0 for daemon and ops jobs
	}

	for _, task := range j.Tasks {
		task.Normalize()
	}
}

// Copy returns a deep copy of the Job. It is expected that callers use recover.
// This job can panic if the deep copy failed as it uses reflection.
func (j *Job) Copy() *Job {
	if j == nil {
		return nil
	}
	nj := new(Job)
	*nj = *j
	nj.Constraints = CopySlice[*LabelSelectorRequirement](nj.Constraints)

	if j.Tasks != nil {
		tasks := make([]*Task, len(nj.Tasks))
		for i, t := range nj.Tasks {
			tasks[i] = t.Copy()
		}
		nj.Tasks = tasks
	}

	nj.Meta = maps.Clone(nj.Meta)
	return nj
}

// Validate is used to check a job for reasonable configuration
func (j *Job) Validate() error {
	errs := errors.Join(
		// Validate the job ID
		validate.NotBlank(j.ID, "missing job ID"),
		validate.NoSpaces(j.ID, "job ID contains a space"),
		validate.NoNullChars(j.ID, "job ID contains a null character"),
		// Validate the job name
		validate.NotBlank(j.Name, "missing job name"),
		validate.NoNullChars(j.Name, "job Name contains a null character"),
		// Validate the job namespace
		validate.NotBlank(j.Namespace, "job must be in a namespace"),
		// More validations
		j.ValidateSubmission(),
	)

	// Validate the task group
	for _, task := range j.Tasks {
		if err := task.Validate(); err != nil {
			outer := fmt.Errorf("task %s validation failed: %v", task.Name, err)
			errs = errors.Join(errs, outer)
		}
	}
	return errs
}

// ValidateSubmission is used to check a job for reasonable configuration when it is submitted.
// It is a subset of Validate that does not check fields with defaults, such as job ID
func (j *Job) ValidateSubmission() error {
	if j == nil {
		return errors.New("empty/nil job")
	}

	var mErr error

	// TODO: Bring it back when we enforce Job Names syntax
	//// Validate job name
	//if j.Name == "" {
	//	mErr = errors.Join(mErr, errors.New("job name cannot be empty"))
	//} else {
	//	if len(j.Name) > 255 {
	//		mErr = errors.Join(mErr, errors.New("job name cannot exceed 255 characters"))
	//	}
	//
	//	// Check that name only contains alphanumeric characters and dashes
	//	validName := regexp.MustCompile(`^[a-zA-Z0-9-]+$`).MatchString
	//	if !validName(j.Name) {
	//		mErr = errors.Join(mErr, errors.New("job name can only contain alphanumeric characters and dashes"))
	//	}
	//}

	switch j.Type {
	case JobTypeService, JobTypeBatch, JobTypeDaemon, JobTypeOps:
	case "":
		mErr = errors.Join(mErr, errors.New("missing job type"))
	default:
		mErr = errors.Join(mErr, fmt.Errorf("invalid job type: %q", j.Type))
	}

	if j.Count < 0 {
		mErr = errors.Join(mErr, errors.New("job count must be >= 0"))
	}
	if j.Count > 1 && (j.Type == JobTypeDaemon || j.Type == JobTypeOps) {
		mErr = errors.Join(mErr, fmt.Errorf("%s jobs cannot specify count > 1 as they run on all matching nodes", j.Type))
	}
	if len(j.Tasks) == 0 {
		mErr = errors.Join(mErr, errors.New("missing job tasks"))
	}
	for idx, constr := range j.Constraints {
		if err := constr.Validate(); err != nil {
			outer := fmt.Errorf("constraint %d validation failed: %s", idx+1, err)
			mErr = errors.Join(mErr, outer)
		}
	}

	// Validate the task group
	for _, task := range j.Tasks {
		if err := task.ValidateSubmission(); err != nil {
			outer := fmt.Errorf("task %s validation failed: %v", task.Name, err)
			mErr = errors.Join(mErr, outer)
		}
	}

	return mErr
}

// SanitizeSubmission is used to sanitize a job for reasonable configuration when it is submitted.
func (j *Job) SanitizeSubmission() (warnings []string) {
	if !j.State.StateType.IsUndefined() {
		warnings = append(warnings, "job state is ignored when submitting a job")
		j.State = NewJobState(JobStateTypeUndefined)
	}
	if j.Revision != 0 {
		warnings = append(warnings, "job revision is ignored when submitting a job")
		j.Revision = 0
	}
	if j.Version != 0 {
		warnings = append(warnings, "job version is ignored when submitting a job")
		j.Version = 0
	}
	if j.CreateTime != 0 {
		warnings = append(warnings, "job create time is ignored when submitting a job")
		j.CreateTime = 0
	}
	if j.ModifyTime != 0 {
		warnings = append(warnings, "job modify time is ignored when submitting a job")
		j.ModifyTime = 0
	}
	if j.Type == JobTypeBatch || j.Type == JobTypeOps {
		if j.ID != "" {
			warnings = append(warnings, "job ID is ignored when submitting a batch job")
			j.ID = ""
		}
	}
	// TODO: remove this once we have multiple tasks per job
	if len(j.Tasks) > 1 {
		warnings = append(warnings, "only one task is supported per job")
		j.Tasks = j.Tasks[:1]
	}
	for k := range j.Meta {
		if strings.HasPrefix(k, MetaReservedPrefix) {
			warnings = append(warnings, fmt.Sprintf("job meta key %q is reserved and will be ignored", k))
			delete(j.Meta, k)
		}
	}
	return warnings
}

// IsTerminal returns true if the job is in a terminal state
func (j *Job) IsTerminal() bool {
	return j.State.StateType.IsTerminal()
}

// IsRerunnable returns true if the job in a state to be re-run
func (j *Job) IsRerunnable() bool {
	return j.State.StateType.IsRerunnable()
}

// Task returns the job task
// TODO: remove this once we have multiple tasks per job
func (j *Job) Task() *Task {
	if j == nil {
		return nil
	}
	if len(j.Tasks) == 0 {
		return nil
	}
	return j.Tasks[0]
}

// GetCreateTime returns the creation time
func (j *Job) GetCreateTime() time.Time {
	return time.Unix(0, j.CreateTime).UTC()
}

// GetModifyTime returns the modify time
func (j *Job) GetModifyTime() time.Time {
	return time.Unix(0, j.ModifyTime).UTC()
}

// AllStorageTypes returns keys of all storage types required by the job
func (j *Job) AllStorageTypes() []string {
	var storageTypes []string
	if j == nil {
		return storageTypes
	}
	for _, task := range j.Tasks {
		storageTypes = append(storageTypes, task.AllStorageTypes()...)
	}
	return storageTypes
}

// IsLongRunning returns true if the job is long running
func (j *Job) IsLongRunning() bool {
	return j.Type == JobTypeService || j.Type == JobTypeDaemon
}

// IsExpired returns true if the job is still running beyond the expiration time
func (j *Job) IsExpired(expirationTime time.Time) bool {
	return !j.IsTerminal() &&
		j.Task().Timeouts.TotalTimeout > 0 &&
		j.GetCreateTime().Before(expirationTime)
}

// OrchestratorID returns the orchestrator ID for the job from its metadata
func (j *Job) OrchestratorID() string {
	if j == nil {
		return ""
	}
	if orchestratorID, ok := j.Meta[MetaOrchestratorID]; ok {
		return orchestratorID
	}
	return j.Meta[MetaOrchestratorIDLegacy]
}

// OrchestrationProtocol returns the orchestrator protocol for the job
func (j *Job) OrchestrationProtocol() Protocol {
	// TODO: Remove this once all jobs have the protocol set when v1.5 is no longer supported
	defaultProtocol := ProtocolBProtocolV2 // Default to legacy protocol
	if j == nil || j.Meta == nil {
		return defaultProtocol
	}
	protocol, ok := j.Meta[MetaOrchestratorProtocol]
	if !ok {
		return defaultProtocol
	}
	return Protocol(protocol)
}

// UnmarshalJSON implements custom JSON unmarshaling to handle
// the Count field differently based on job type
// This allows us to:
// 1. Set default Count=1 for batch and service jobs when omitted in JSON
// 2. Preserve explicit Count=0 when specified (used to stop all executions)
func (j *Job) UnmarshalJSON(data []byte) error {
	type alias Job // Create alias to avoid recursion
	type job struct {
		*alias
		Count *int `json:"Count,omitempty"`
	}
	aux := &job{
		alias: (*alias)(j),
	}
	if err := json.Unmarshal(data, aux); err != nil {
		return err
	}

	// Handle Count field based on job type
	switch j.Type {
	case JobTypeBatch, JobTypeService:
		if aux.Count == nil {
			j.Count = 1 // Default to 1 if not present in JSON
		} else {
			j.Count = *aux.Count // Use explicitly set value
		}
	default:
		if aux.Count != nil {
			j.Count = *aux.Count
		}
	}
	return nil
}

// CompareWith returns a diff string between this job and another job.
// The comparison ignores fields that are not relevant to job specification equivalence
// such as ID, State, Version, Revision, CreateTime, and ModifyTime.
// It also ignores internal metadata entries with keys starting with the reserved prefix,
// and performs case-insensitive comparison of resource values.
// Returns an empty string if jobs are equivalent, otherwise returns a human-readable diff.
func (j *Job) CompareWith(otherJob *Job) string {
	diffResult := cmp.Diff(j, otherJob,
		// Ignoring these fields since they are not part of the comparison
		cmpopts.IgnoreFields(
			Job{}, "ID", "State", "Version", "Revision", "CreateTime", "ModifyTime",
		),
		cmp.FilterPath(func(p cmp.Path) bool {
			// Check if we're looking at a Meta map entry
			if len(p) >= 2 && p.Index(-2).String() == ".Meta" {
				// Get the map key
				if mapKey, ok := p.Index(-1).(cmp.MapIndex); ok {
					// Check if the key starts with "bacalhau.org" which is reserved
					// for internal bacalhau use, thus no need to compare.
					if key, ok := mapKey.Key().Interface().(string); ok {
						return strings.HasPrefix(key, MetaReservedPrefix)
					}
				}
			}
			return false
		}, cmp.Ignore()),
		// Case-insensitive comparison for ResourcesConfig string fields. These are
		// treated as case-insensitive due units used.
		cmp.Transformer("ResourcesConfigCaseInsensitive", func(r *ResourcesConfig) *ResourcesConfig {
			if r == nil {
				return nil
			}
			result := r.Copy()
			result.CPU = strings.ToLower(result.CPU)
			result.Memory = strings.ToLower(result.Memory)
			result.Disk = strings.ToLower(result.Disk)
			result.GPU = strings.ToLower(result.GPU)
			return result
		}),
	)

	return diffResult
}
