package model

import (
	"errors"

	"github.com/MAX-API-Next/MAX-API/common"

	"gorm.io/gorm"
)

type SystemTaskStatus string

const (
	SystemTaskStatusPending   SystemTaskStatus = "pending"
	SystemTaskStatusRunning   SystemTaskStatus = "running"
	SystemTaskStatusSucceeded SystemTaskStatus = "succeeded"
	SystemTaskStatusFailed    SystemTaskStatus = "failed"

	SystemTaskTypeLogCleanup     = "log_cleanup"
	SystemTaskTypeChannelTest    = "channel_test"
	SystemTaskTypeModelUpdate    = "model_update"
	SystemTaskTypeMidjourneyPoll = "midjourney_poll"
	SystemTaskTypeAsyncTaskPoll  = "async_task_poll"
)

var ErrSystemTaskLockLost = errors.New("system task lock lost")

type SystemTask struct {
	ID          int64            `json:"id" gorm:"primaryKey"`
	TaskID      string           `json:"task_id" gorm:"type:varchar(64);uniqueIndex"`
	Type        string           `json:"type" gorm:"type:varchar(64);index"`
	Status      SystemTaskStatus `json:"status" gorm:"type:varchar(32);index"`
	ActiveKey   *string          `json:"active_key,omitempty" gorm:"type:varchar(64);uniqueIndex"`
	Payload     string           `json:"payload" gorm:"type:text"`
	State       string           `json:"state" gorm:"type:text"`
	Result      string           `json:"result" gorm:"type:text"`
	Error       string           `json:"error" gorm:"type:text"`
	LockedBy    string           `json:"locked_by" gorm:"type:varchar(128);index"`
	LockedUntil int64            `json:"locked_until" gorm:"bigint;index"`
	CreatedAt   int64            `json:"created_at" gorm:"bigint;index"`
	UpdatedAt   int64            `json:"updated_at" gorm:"bigint;index"`
}

type SystemTaskResponse struct {
	ID          int64            `json:"id"`
	TaskID      string           `json:"task_id"`
	Type        string           `json:"type"`
	Status      SystemTaskStatus `json:"status"`
	ActiveKey   string           `json:"active_key,omitempty"`
	Payload     any              `json:"payload"`
	State       any              `json:"state"`
	Result      any              `json:"result"`
	Error       string           `json:"error"`
	LockedBy    string           `json:"locked_by"`
	LockedUntil int64            `json:"locked_until"`
	CreatedAt   int64            `json:"created_at"`
	UpdatedAt   int64            `json:"updated_at"`
}

func (task *SystemTask) BeforeCreate(_ *gorm.DB) error {
	now := common.GetTimestamp()
	if task.CreatedAt == 0 {
		task.CreatedAt = now
	}
	if task.UpdatedAt == 0 {
		task.UpdatedAt = now
	}
	return nil
}

func GenerateSystemTaskID() (string, error) {
	key, err := common.GenerateRandomCharsKey(32)
	if err != nil {
		return "", err
	}
	return "systask_" + key, nil
}

func CreateSystemTask(taskType string, activeKey string, payload any, state any) (*SystemTask, error) {
	taskID, err := GenerateSystemTaskID()
	if err != nil {
		return nil, err
	}
	payloadText, err := marshalSystemTaskJSON(payload)
	if err != nil {
		return nil, err
	}
	stateText, err := marshalSystemTaskJSON(state)
	if err != nil {
		return nil, err
	}

	task := &SystemTask{
		TaskID:  taskID,
		Type:    taskType,
		Status:  SystemTaskStatusPending,
		Payload: payloadText,
		State:   stateText,
	}
	if activeKey != "" {
		task.ActiveKey = &activeKey
	}

	if err := DB.Create(task).Error; err != nil {
		return nil, err
	}
	return task, nil
}

func GetSystemTaskByTaskID(taskID string) (*SystemTask, error) {
	var task SystemTask
	if err := DB.Where("task_id = ?", taskID).First(&task).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, nil
		}
		return nil, err
	}
	return &task, nil
}

func GetActiveSystemTask(taskType string) (*SystemTask, error) {
	var task SystemTask
	err := DB.Where("type = ? AND active_key IS NOT NULL", taskType).
		Where("status IN ?", activeSystemTaskStatuses()).
		Order("id desc").
		First(&task).Error
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, nil
		}
		return nil, err
	}
	return &task, nil
}

func FindRunnableSystemTasks(taskType string, now int64, limit int) ([]*SystemTask, error) {
	var tasks []*SystemTask
	if limit <= 0 {
		limit = 1
	}
	err := DB.Where("type = ? AND status IN ? AND (locked_until = 0 OR locked_until < ?)", taskType, activeSystemTaskStatuses(), now).
		Order("id asc").
		Limit(limit).
		Find(&tasks).Error
	return tasks, err
}

func ListSystemTasks(limit int) ([]*SystemTask, error) {
	if limit <= 0 {
		limit = 20
	}
	if limit > 100 {
		limit = 100
	}
	var tasks []*SystemTask
	err := DB.Order("id desc").Limit(limit).Find(&tasks).Error
	return tasks, err
}

func ClaimSystemTask(id int64, taskType string, runnerID string, lockUntil int64) (*SystemTask, bool, error) {
	now := common.GetTimestamp()
	result := DB.Model(&SystemTask{}).
		Where("id = ? AND type = ? AND status IN ? AND (locked_until = 0 OR locked_until < ? OR locked_by = ?)", id, taskType, activeSystemTaskStatuses(), now, runnerID).
		Updates(map[string]any{
			"status":       SystemTaskStatusRunning,
			"locked_by":    runnerID,
			"locked_until": lockUntil,
			"updated_at":   now,
		})
	if result.Error != nil {
		return nil, false, result.Error
	}
	if result.RowsAffected == 0 {
		return nil, false, nil
	}

	var task SystemTask
	if err := DB.Where("id = ?", id).First(&task).Error; err != nil {
		return nil, false, err
	}
	return &task, true, nil
}

func UpdateSystemTaskState(taskID string, lockedBy string, state any, lockUntil int64) error {
	stateText, err := marshalSystemTaskJSON(state)
	if err != nil {
		return err
	}
	result := DB.Model(&SystemTask{}).
		Where("task_id = ? AND status = ? AND locked_by = ?", taskID, SystemTaskStatusRunning, lockedBy).
		Updates(map[string]any{
			"state":        stateText,
			"locked_until": lockUntil,
			"updated_at":   common.GetTimestamp(),
		})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return ErrSystemTaskLockLost
	}
	return nil
}

func FinishSystemTask(taskID string, lockedBy string, status SystemTaskStatus, resultPayload any, errorMessage string) error {
	resultText, err := marshalSystemTaskJSON(resultPayload)
	if err != nil {
		return err
	}
	result := DB.Model(&SystemTask{}).
		Where("task_id = ? AND status = ? AND locked_by = ?", taskID, SystemTaskStatusRunning, lockedBy).
		Updates(map[string]any{
			"status":       status,
			"active_key":   nil,
			"result":       resultText,
			"error":        errorMessage,
			"locked_by":    "",
			"locked_until": 0,
			"updated_at":   common.GetTimestamp(),
		})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return ErrSystemTaskLockLost
	}
	return nil
}

func (task *SystemTask) DecodePayload(v any) error {
	return decodeSystemTaskJSONString(task.Payload, v)
}

func (task *SystemTask) DecodeState(v any) error {
	return decodeSystemTaskJSONString(task.State, v)
}

func (task *SystemTask) ToResponse() SystemTaskResponse {
	activeKey := ""
	if task.ActiveKey != nil {
		activeKey = *task.ActiveKey
	}
	return SystemTaskResponse{
		ID:          task.ID,
		TaskID:      task.TaskID,
		Type:        task.Type,
		Status:      task.Status,
		ActiveKey:   activeKey,
		Payload:     decodeSystemTaskJSONValue(task.Payload),
		State:       decodeSystemTaskJSONValue(task.State),
		Result:      decodeSystemTaskJSONValue(task.Result),
		Error:       task.Error,
		LockedBy:    task.LockedBy,
		LockedUntil: task.LockedUntil,
		CreatedAt:   task.CreatedAt,
		UpdatedAt:   task.UpdatedAt,
	}
}

func activeSystemTaskStatuses() []string {
	return []string{string(SystemTaskStatusPending), string(SystemTaskStatusRunning)}
}

func marshalSystemTaskJSON(v any) (string, error) {
	if v == nil {
		return "", nil
	}
	data, err := common.Marshal(v)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func decodeSystemTaskJSONString(data string, v any) error {
	if data == "" {
		return nil
	}
	return common.UnmarshalJsonStr(data, v)
}

func decodeSystemTaskJSONValue(data string) any {
	if data == "" {
		return nil
	}
	var value any
	if err := common.UnmarshalJsonStr(data, &value); err != nil {
		return data
	}
	return value
}
