package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"github.com/riverqueue/river"

	"github.com/nomadcode/nomadcode-core/internal/agent"
	"github.com/nomadcode/nomadcode-core/internal/authoring"
	"github.com/nomadcode/nomadcode-core/internal/model"
	"github.com/nomadcode/nomadcode-core/internal/notification"
	"github.com/nomadcode/nomadcode-core/internal/storage"
	"github.com/nomadcode/nomadcode-core/internal/workflow"
)

type TaskJobArgs struct {
	TaskID string `json:"task_id"`
}

func (TaskJobArgs) Kind() string {
	return "task_run"
}

// TaskLifecycle handles task state transitions and canonical metadata rules.
//
// Responsibility boundaries:
//   - Workflow service owns create/list/get/enqueue API semantics.
//   - Lifecycle owns all canonical status writes and metadata merge.
//   - Scheduler worker owns invoking StartTask, running execution, then invoking CompleteTask or FailTask.
//   - Provider projection failure must not roll back canonical task state.
type TaskLifecycle interface {
	StartTask(ctx context.Context, id string) (storage.Task, error)
	CompleteTask(ctx context.Context, id string, result json.RawMessage) (storage.Task, error)
	FailTask(ctx context.Context, id string, message string) (storage.Task, error)
	FailTaskWithMetadata(ctx context.Context, id string, input workflow.FailureInput) (storage.Task, error)
	MergeTaskMetadata(ctx context.Context, id string, updates map[string]any) (storage.Task, error)
}

type TaskWorker struct {
	river.WorkerDefaults[TaskJobArgs]

	Lifecycle     TaskLifecycle
	Notifications *notification.Service
	Agent         agent.Client
	Model         model.Client
	RunTimeout    time.Duration
	Logger        *slog.Logger
}

func (w *TaskWorker) Work(ctx context.Context, job *river.Job[TaskJobArgs]) error {
	taskID := job.Args.TaskID
	if w.Logger != nil {
		w.Logger.Info("task job started", "task_id", taskID)
	}

	task, err := w.Lifecycle.StartTask(ctx, taskID)
	if err != nil {
		return err
	}

	// Record authoring_run_state=in_progress before execution so a stale
	// detection layer can compare authoring_run_updated_at against a threshold
	// without waiting for the run to complete or fail.
	if _, isAuthoring := authoring.BuildAuthoringGenerateInput(task); isAuthoring {
		if _, mergeErr := w.Lifecycle.MergeTaskMetadata(ctx, taskID, map[string]any{
			workflow.MetadataKeyAuthoringRunState:     "in_progress",
			workflow.MetadataKeyAuthoringRunUpdatedAt: time.Now().UTC().Format(time.RFC3339),
		}); mergeErr != nil && w.Logger != nil {
			w.Logger.Warn("authoring in_progress metadata merge failed", "task_id", taskID, "error", mergeErr)
		}
	}

	attempt, reason, occurredAt := parseTaskEventContext(task, notification.TaskEventRunning)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventRunning,
		TaskID:     task.ID,
		Title:      task.Title,
		Status:     task.Status,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})

	runCtx := ctx
	if w.RunTimeout > 0 {
		var cancel context.CancelFunc
		runCtx, cancel = context.WithTimeout(ctx, w.RunTimeout)
		defer cancel()
	}

	result, message, err := w.runTask(runCtx, task)
	if err != nil {
		w.markFailedWithTask(taskID, task, err)
		return err
	}

	// Plane-origin authoring task의 경우, bridge 성공만으로는 task를 완료하지 않고
	// develop match를 대기하는 상태(running)를 유지한다.
	_, isAuthoring := authoring.BuildAuthoringGenerateInput(task)
	if isAuthoring {
		var resMap map[string]any
		if err := json.Unmarshal(result, &resMap); err == nil {
			if state, ok := resMap[workflow.MetadataKeyAuthoringRunState].(string); ok && state == "in_progress" {
				updates := map[string]any{
					workflow.MetadataKeyAuthoringRunState:     "in_progress",
					workflow.MetadataKeyStatusReason:          message,
					workflow.MetadataKeyWaitType:              "develop_match",
					workflow.MetadataKeyAuthoringRunUpdatedAt: time.Now().UTC().Format(time.RFC3339),
				}
				for k, v := range resMap {
					updates[k] = v
				}

				if _, mergeErr := w.Lifecycle.MergeTaskMetadata(ctx, taskID, updates); mergeErr != nil && w.Logger != nil {
					w.Logger.Warn("authoring running metadata merge failed", "task_id", taskID, "error", mergeErr)
				}

				if w.Logger != nil {
					w.Logger.Info("task job remains running for develop match", "task_id", taskID)
				}
				return nil
			}
		}
	}

	task, err = w.Lifecycle.CompleteTask(ctx, taskID, result)
	if err != nil {
		w.markFailed(taskID, err)
		return err
	}

	attempt, reason, occurredAt = parseTaskEventContext(task, notification.TaskEventCompleted)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventCompleted,
		TaskID:     task.ID,
		Title:      task.Title,
		Status:     task.Status,
		Message:    message,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})

	if w.Logger != nil {
		w.Logger.Info("task job completed", "task_id", taskID)
	}
	return nil
}

func (w *TaskWorker) runTask(ctx context.Context, task storage.Task) (json.RawMessage, string, error) {
	// Plane-origin authoring tasks always route through the IOP OpenAI-compatible
	// HTTP bridge even when an A2A agent client is configured, because IOP runs
	// workspace-bound agent authoring rather than a generic A2A task.
	generateInput, isAuthoring := authoring.BuildAuthoringGenerateInput(task)
	if isAuthoring {
		if w.Model == nil {
			return nil, "", fmt.Errorf("model client is required for authoring tasks")
		}
		generated, err := w.Model.Generate(ctx, generateInput)
		if err != nil {
			return nil, "", err
		}
		decision := authoring.DecideAuthoringResult(authoring.DecisionInput{
			BridgeSuccess: true,
		})
		result := map[string]any{
			"message":                  generated.Text,
			"mode":                     "authoring_run",
			"model":                    generated.Model,
			"response_id":              generated.ID,
			"authoring_run_state":      decision.State,
			"authoring_run_updated_at": time.Now().UTC().Format(time.RFC3339),
			"usage": map[string]int{
				"input_tokens":  generated.Usage.InputTokens,
				"output_tokens": generated.Usage.OutputTokens,
				"total_tokens":  generated.Usage.TotalTokens,
			},
		}
		raw, err := json.Marshal(result)
		if err != nil {
			return nil, "", err
		}
		return raw, decision.Reason, nil
	}

	if w.Agent != nil {
		return w.runAgentTask(ctx, task)
	}

	if w.Model == nil {
		select {
		case <-time.After(500 * time.Millisecond):
		case <-ctx.Done():
			return nil, "", ctx.Err()
		}
		return json.RawMessage(`{"message":"dummy task completed","mode":"dummy"}`), "dummy task completed", nil
	}

	generated, err := w.Model.Generate(ctx, buildGenerateInput(task))
	if err != nil {
		return nil, "", err
	}

	result := map[string]any{
		"message":     generated.Text,
		"mode":        "openai_responses",
		"model":       generated.Model,
		"response_id": generated.ID,
		"usage": map[string]int{
			"input_tokens":  generated.Usage.InputTokens,
			"output_tokens": generated.Usage.OutputTokens,
			"total_tokens":  generated.Usage.TotalTokens,
		},
	}
	raw, err := json.Marshal(result)
	if err != nil {
		return nil, "", err
	}
	return raw, "model task completed", nil
}

func (w *TaskWorker) runAgentTask(ctx context.Context, task storage.Task) (json.RawMessage, string, error) {
	result, err := w.Agent.SendMessage(ctx, buildAgentInput(task))
	if err != nil {
		return nil, "", err
	}

	output := map[string]any{
		"mode": "a2a",
	}
	if len(result.Raw) > 0 {
		output["raw"] = json.RawMessage(result.Raw)
	}

	message := "a2a task completed"
	if result.Message != nil {
		text := messageText(*result.Message)
		if text != "" {
			message = text
			output["message"] = text
		}
		output["message_id"] = result.Message.MessageID
		raw, err := json.Marshal(output)
		if err != nil {
			return nil, "", err
		}
		return raw, message, nil
	}

	if result.Task == nil {
		return nil, "", fmt.Errorf("a2a result does not include message or task")
	}

	state := result.Task.Status.State
	output["a2a_task_id"] = result.Task.ID
	output["a2a_state"] = state

	text := taskOutputText(*result.Task)
	if text != "" {
		message = text
		output["message"] = text
	}

	switch state {
	case agent.TaskStateCompleted:
	case agent.TaskStateFailed, agent.TaskStateCanceled, agent.TaskStateRejected, agent.TaskStateInputNeeded, agent.TaskStateAuthRequired:
		if text == "" {
			text = fmt.Sprintf("a2a task ended with state %s", state)
		}
		return nil, "", fmt.Errorf("%s", text)
	default:
		return nil, "", fmt.Errorf("a2a blocking call returned non-terminal task state %s", state)
	}

	raw, err := json.Marshal(output)
	if err != nil {
		return nil, "", err
	}
	return raw, message, nil
}

func buildGenerateInput(task storage.Task) model.GenerateInput {
	input := strings.TrimSpace(task.Title)
	instructions := ""

	if len(task.Payload) > 0 && string(task.Payload) != "null" {
		if payloadInput, payloadInstructions, ok := parsePayloadPrompt(task.Payload); ok {
			if payloadInput != "" {
				input = payloadInput
			}
			instructions = payloadInstructions
		} else if string(task.Payload) != "{}" {
			input = fmt.Sprintf("Task title: %s\nPayload:\n%s", task.Title, string(task.Payload))
		}
	}

	return model.GenerateInput{
		Input:        input,
		Instructions: instructions,
		Metadata: map[string]string{
			"task_id": task.ID,
			"source":  task.Source,
		},
	}
}

func buildAgentInput(task storage.Task) agent.SendMessageInput {
	generateInput := buildGenerateInput(task)
	text := strings.TrimSpace(generateInput.Input)
	if generateInput.Instructions != "" {
		text = fmt.Sprintf("Instructions:\n%s\n\nTask:\n%s", generateInput.Instructions, text)
	}
	if text == "" {
		text = task.Title
	}

	return agent.SendMessageInput{
		Text:                text,
		AcceptedOutputModes: []string{"text/plain", "application/json"},
		Blocking:            true,
		Metadata: map[string]any{
			"nomadcode_task_id": task.ID,
			"source":            task.Source,
			"title":             task.Title,
		},
	}
}

func parsePayloadPrompt(payload json.RawMessage) (input string, instructions string, ok bool) {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(payload, &fields); err != nil {
		return "", "", false
	}
	ok = true
	input = firstStringField(fields, "prompt", "message", "input")
	instructions = firstStringField(fields, "instructions", "system")
	return input, instructions, ok
}

func firstStringField(fields map[string]json.RawMessage, names ...string) string {
	for _, name := range names {
		raw, exists := fields[name]
		if !exists {
			continue
		}
		var value string
		if err := json.Unmarshal(raw, &value); err == nil {
			if value = strings.TrimSpace(value); value != "" {
				return value
			}
		}
	}
	return ""
}

func taskOutputText(task agent.Task) string {
	if task.Status.Message != nil {
		if text := messageText(*task.Status.Message); text != "" {
			return text
		}
	}

	for _, artifact := range task.Artifacts {
		if text := partsText(artifact.Parts); text != "" {
			return text
		}
	}

	for i := len(task.History) - 1; i >= 0; i-- {
		if task.History[i].Role == "agent" || task.History[i].Role == "assistant" {
			if text := messageText(task.History[i]); text != "" {
				return text
			}
		}
	}
	for i := len(task.History) - 1; i >= 0; i-- {
		if text := messageText(task.History[i]); text != "" {
			return text
		}
	}

	return ""
}

func messageText(message agent.Message) string {
	return partsText(message.Parts)
}

func partsText(parts []agent.Part) string {
	texts := make([]string, 0, len(parts))
	for _, part := range parts {
		text := strings.TrimSpace(part.Text)
		if text != "" {
			texts = append(texts, text)
		}
	}
	return strings.Join(texts, "\n")
}

func (w *TaskWorker) notifyTaskEvent(ctx context.Context, event notification.TaskEvent) {
	if w.Notifications == nil {
		return
	}
	err := w.Notifications.NotifyTaskEvent(ctx, event)
	if err != nil && w.Logger != nil {
		w.Logger.Warn("task notification failed", "task_id", event.TaskID, "type", event.Type, "error", err)
	}
}

// markFailedWithTask is markFailed with the originating task available so
// authoring-specific failure metadata can be appended before the lifecycle call.
func (w *TaskWorker) markFailedWithTask(taskID string, task storage.Task, err error) {
	if err == nil || w.Lifecycle == nil {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	msg := err.Error()
	failureType := workflow.FailureTypeExecution
	if errors.Is(err, context.DeadlineExceeded) {
		msg = "timeout"
		failureType = workflow.FailureTypeTimeout
	}

	failInput := workflow.FailureInput{Message: msg, Type: failureType}

	if _, isAuthoring := authoring.BuildAuthoringGenerateInput(task); isAuthoring {
		decision := authoring.DecideAuthoringResult(authoring.DecisionInput{BridgeSuccess: false})
		failInput.ExtraMetadata = map[string]any{
			workflow.MetadataKeyAuthoringRunState:        decision.State,
			workflow.MetadataKeyAuthoringRunUpdatedAt:    time.Now().UTC().Format(time.RFC3339),
			workflow.MetadataKeyAuthoringFailureType:     string(failureType),
			workflow.MetadataKeyAuthoringFailureCategory: decision.FailureCategory,
		}
	}

	failed, failErr := w.Lifecycle.FailTaskWithMetadata(ctx, taskID, failInput)
	if failErr != nil {
		if w.Logger != nil {
			w.Logger.Error("failed to mark task failed", "task_id", taskID, "error", failErr)
		}
		return
	}

	attempt, reason, occurredAt := parseTaskEventContext(failed, notification.TaskEventFailed)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventFailed,
		TaskID:     failed.ID,
		Title:      failed.Title,
		Status:     failed.Status,
		Message:    msg,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})
}

func (w *TaskWorker) markFailed(taskID string, err error) {
	if err == nil || w.Lifecycle == nil {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	msg := err.Error()
	failureType := workflow.FailureTypeExecution
	if errors.Is(err, context.DeadlineExceeded) {
		msg = "timeout"
		failureType = workflow.FailureTypeTimeout
	}

	task, failErr := w.Lifecycle.FailTaskWithMetadata(ctx, taskID, workflow.FailureInput{
		Message: msg,
		Type:    failureType,
	})
	if failErr != nil {
		if w.Logger != nil {
			w.Logger.Error("failed to mark task failed", "task_id", taskID, "error", failErr)
		}
		return
	}

	attempt, reason, occurredAt := parseTaskEventContext(task, notification.TaskEventFailed)
	w.notifyTaskEvent(ctx, notification.TaskEvent{
		Type:       notification.TaskEventFailed,
		TaskID:     task.ID,
		Title:      task.Title,
		Status:     task.Status,
		Message:    msg,
		Attempt:    attempt,
		Reason:     reason,
		OccurredAt: occurredAt,
	})
}

func parseTaskEventContext(task storage.Task, eventType notification.TaskEventType) (int, string, time.Time) {
	attempt := 0
	reason := ""
	occurredAt := task.UpdatedAt

	if len(task.Metadata) > 0 && string(task.Metadata) != "null" {
		var meta map[string]any
		if err := json.Unmarshal(task.Metadata, &meta); err == nil {
			if val, ok := meta[workflow.MetadataKeyAttempt]; ok {
				switch v := val.(type) {
				case float64:
					attempt = int(v)
				case int:
					attempt = v
				case int64:
					attempt = int(v)
				}
			}
			if eventType == notification.TaskEventFailed {
				if val, ok := meta[workflow.MetadataKeyStatusReason]; ok {
					if str, ok := val.(string); ok {
						reason = str
					}
				}
			}
		}
	}
	return attempt, reason, occurredAt
}