// Package errtypes provides custom error types for fetch_ml
package errtypes

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

// DataFetchError represents an error that occurred while fetching a dataset
// from the NAS to the ML server.
type DataFetchError struct {
	Err     error
	Dataset string
	JobName string
}

func (e *DataFetchError) Error() string {
	return fmt.Sprintf("failed to fetch dataset %s for job %s: %v",
		e.Dataset, e.JobName, e.Err)
}
func (e *DataFetchError) Unwrap() error {
	return e.Err
}

// TaskExecutionError represents an error during task execution.
type TaskExecutionError struct {
	Timestamp   time.Time         `json:"timestamp"`
	Err         error             `json:"-"`
	Context     map[string]string `json:"context,omitempty"`
	TaskID      string            `json:"task_id"`
	JobName     string            `json:"job_name"`
	Phase       string            `json:"phase"`
	Message     string            `json:"message"`
	Recoverable bool              `json:"recoverable"`
}

// Error returns the error message.
func (e *TaskExecutionError) Error() string {
	if e.Message != "" {
		return fmt.Sprintf("task %s (%s) failed during %s: %s",
			e.TaskID[:8], e.JobName, e.Phase, e.Message)
	}
	return fmt.Sprintf("task %s (%s) failed during %s: %v",
		e.TaskID[:8], e.JobName, e.Phase, e.Err)
}

// Unwrap returns the underlying error.
func (e *TaskExecutionError) Unwrap() error {
	return e.Err
}

// MarshalJSON returns a JSON representation of the error.
func (e *TaskExecutionError) MarshalJSON() ([]byte, error) {
	type Alias TaskExecutionError
	return json.Marshal(&struct {
		*Alias
		Error string `json:"error,omitempty"`
	}{
		Alias: (*Alias)(e),
		Error: func() string {
			if e.Err != nil {
				return e.Err.Error()
			}
			return ""
		}(),
	})
}

// IsRecoverable returns true if the error is retryable.
func (e *TaskExecutionError) IsRecoverable() bool {
	return e.Recoverable
}

// WithContext adds context to the error.
func (e *TaskExecutionError) WithContext(key, value string) *TaskExecutionError {
	if e.Context == nil {
		e.Context = make(map[string]string)
	}
	e.Context[key] = value
	return e
}
