package runtime

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

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

// Worker owns runner-facing state transitions. Execution is owned by runners
// such as ServerRunner or runner-compatible clients.
type Worker struct {
	store *automation.Store
	now   func() time.Time
}

func NewWorker(store *automation.Store, now func() time.Time, _ any) *Worker {
	return &Worker{store: store, now: now}
}

func (w *Worker) AcceptLeasedJob(ctx context.Context, job automation.AutomationJob, leaseOwner string) (automation.AutomationJobRun, error) {
	if w == nil || w.store == nil {
		return automation.AutomationJobRun{}, fmt.Errorf("automation worker store is nil")
	}
	now := w.currentTime()
	jobRun, err := w.store.GetJobRun(ctx, job.WorkspaceID, job.JobRunID)
	if err != nil {
		return automation.AutomationJobRun{}, err
	}
	if jobRun.Status == automation.AutomationJobRunStatusRunning && jobRun.LeaseOwner == leaseOwner {
		return jobRun, nil
	}
	jobRun.Status = automation.AutomationJobRunStatusRunning
	jobRun.Attempt = job.Attempt
	jobRun.LeaseOwner = leaseOwner
	jobRun.LeaseUntil = job.LeaseUntil
	if jobRun.StartedAt == nil {
		started := now
		jobRun.StartedAt = &started
	}
	jobRun.UpdatedAt = now
	if err := w.store.MarkJobAssignmentRunning(ctx, job.ID, leaseOwner, jobRun, now); err != nil {
		return automation.AutomationJobRun{}, err
	}
	return jobRun, nil
}

func (w *Worker) CompleteLeasedJob(ctx context.Context, job automation.AutomationJob, leaseOwner string, outputJSON string) error {
	if w == nil || w.store == nil {
		return fmt.Errorf("automation worker store is nil")
	}
	finished := w.currentTime()
	jobRun, err := w.store.GetJobRun(ctx, job.WorkspaceID, job.JobRunID)
	if err != nil {
		return err
	}
	if jobRun.Status == automation.AutomationJobRunStatusSucceeded {
		currentJob, getErr := w.store.GetJob(ctx, job.ID)
		if getErr == nil && currentJob.Status == automation.AutomationJobStatusDone {
			return w.finalizeAfterSuccess(ctx, jobRun.WorkspaceID, jobRun.RunID, finished)
		}
	}
	jobRun.Status = automation.AutomationJobRunStatusSucceeded
	jobRun.Attempt = job.Attempt
	jobRun.OutputJSON = normalizeJSONObject(outputJSON)
	jobRun.CompletedAt = &finished
	jobRun.UpdatedAt = finished
	if err := w.store.CompleteJobAssignmentSuccess(ctx, job.ID, leaseOwner, jobRun, w.jobRunEvent(jobRun, "job_succeeded", "Automation job succeeded", finished), finished); err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) && w.successAlreadyApplied(ctx, job) {
			return w.finalizeAfterSuccess(ctx, jobRun.WorkspaceID, jobRun.RunID, finished)
		}
		return err
	}
	return w.finalizeAfterSuccess(ctx, jobRun.WorkspaceID, jobRun.RunID, finished)
}

func (w *Worker) finalizeAfterSuccess(ctx context.Context, workspaceID string, runID string, finished time.Time) error {
	if _, err := NewScheduler(w.store, w.now).ScheduleRun(ctx, workspaceID, runID); err != nil {
		return err
	}
	_, err := w.store.CompleteRunIfTerminal(ctx, workspaceID, runID, finished)
	return err
}

func (w *Worker) FailLeasedJob(ctx context.Context, job automation.AutomationJob, leaseOwner string, message string, outputJSON string) error {
	return w.failLeasedJob(ctx, job, leaseOwner, message, nil, outputJSON, false)
}

// FailLeasedJobTerminal records a leased job failure without scheduling another attempt.
func (w *Worker) FailLeasedJobTerminal(ctx context.Context, job automation.AutomationJob, leaseOwner string, message string, outputJSON string) error {
	return w.failLeasedJob(ctx, job, leaseOwner, message, nil, outputJSON, true)
}

// FailLeasedJobTerminalWithMetadata records a terminal failure with additional
// stable metadata for operators while preserving the human-readable message.
func (w *Worker) FailLeasedJobTerminalWithMetadata(ctx context.Context, job automation.AutomationJob, leaseOwner string, message string, metadata map[string]any, outputJSON string) error {
	return w.failLeasedJob(ctx, job, leaseOwner, message, metadata, outputJSON, true)
}

func (w *Worker) failLeasedJob(ctx context.Context, job automation.AutomationJob, leaseOwner string, message string, metadata map[string]any, outputJSON string, terminal bool) error {
	if w == nil || w.store == nil {
		return fmt.Errorf("automation worker store is nil")
	}
	finished := w.currentTime()
	jobRun, err := w.store.GetJobRun(ctx, job.WorkspaceID, job.JobRunID)
	if err != nil {
		return err
	}
	if w.failureAlreadyApplied(ctx, job) {
		return nil
	}
	attemptsExhausted := terminal || job.Attempt >= job.MaxAttempts
	jobRun.Status = automation.AutomationJobRunStatusQueued
	eventType := "job_retry_scheduled"
	eventSummary := "Automation job retry scheduled"
	completedAt := (*time.Time)(nil)
	if attemptsExhausted {
		jobRun.Status = automation.AutomationJobRunStatusFailed
		eventType = "job_failed"
		eventSummary = "Automation job failed"
		completedAt = &finished
	}
	jobRun.Attempt = job.Attempt
	jobRun.ErrorJSON = automationFailureMetadataJSON(message, metadata)
	if outputJSON != "" {
		jobRun.OutputJSON = normalizeJSONObject(outputJSON)
	}
	jobRun.CompletedAt = completedAt
	jobRun.UpdatedAt = finished
	retryAt := finished.Add(automation.RetryBackoff(job.Attempt))
	jobRun.ScheduledAt = retryAt
	failure, err := w.store.FailJobAssignment(ctx, job.ID, leaseOwner, jobRun, w.jobRunEvent(jobRun, eventType, eventSummary, finished), retryAt, finished, terminal)
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) && w.failureAlreadyApplied(ctx, job) {
			return nil
		}
		return err
	}
	if failure.AttemptsExhaust {
		_, err := w.store.FailRunAfterJobFailureWithMetadata(ctx, jobRun.WorkspaceID, jobRun.RunID, jobRun.ID, message, metadata, finished)
		return err
	}
	return nil
}

func (w *Worker) CancelRun(ctx context.Context, workspaceID string, runID string, actorPrincipalID string, reason string) error {
	if w == nil || w.store == nil {
		return fmt.Errorf("automation worker store is nil")
	}
	return w.store.CancelRunRuntime(ctx, workspaceID, runID, actorPrincipalID, reason, w.currentTime())
}

func (w *Worker) successAlreadyApplied(ctx context.Context, job automation.AutomationJob) bool {
	currentJob, err := w.store.GetJob(ctx, job.ID)
	if err != nil || currentJob.Status != automation.AutomationJobStatusDone {
		return false
	}
	jobRun, err := w.store.GetJobRun(ctx, job.WorkspaceID, job.JobRunID)
	return err == nil && jobRun.Status == automation.AutomationJobRunStatusSucceeded
}

func (w *Worker) failureAlreadyApplied(ctx context.Context, job automation.AutomationJob) bool {
	currentJob, err := w.store.GetJob(ctx, job.ID)
	if err != nil {
		return false
	}
	currentJobRun, err := w.store.GetJobRun(ctx, job.WorkspaceID, job.JobRunID)
	if err != nil {
		return false
	}
	if currentJob.Status == automation.AutomationJobStatusFailed && currentJobRun.Status == automation.AutomationJobRunStatusFailed {
		return true
	}
	if currentJob.Status == automation.AutomationJobStatusReady && currentJobRun.Status == automation.AutomationJobRunStatusQueued && currentJob.Attempt >= job.Attempt {
		return true
	}
	return false
}

func (w *Worker) jobRunEvent(jobRun automation.AutomationJobRun, eventType string, summary string, now time.Time) automation.AutomationEvent {
	return automation.AutomationEvent{ID: w.eventIDAt(eventType, jobRun.ID, now), WorkspaceID: jobRun.WorkspaceID, RunID: jobRun.RunID, JobRunID: jobRun.ID, EventType: eventType, Summary: summary, CreatedAt: now}
}

func (w *Worker) eventIDAt(eventType string, seed string, now time.Time) string {
	return fmt.Sprintf("wfe_%s_%s_%d", eventType, eventIDSuffix(seed), now.UnixNano())
}

func eventIDSuffix(seed string) string {
	suffix := seed
	if len(suffix) > 16 {
		suffix = suffix[len(suffix)-16:]
	}
	return suffix
}

func (w *Worker) currentTime() time.Time {
	if w.now != nil {
		return w.now()
	}
	return time.Now().UTC()
}

func normalizeJSONObject(value string) string {
	if value == "" {
		return "{}"
	}
	var payload any
	if err := json.Unmarshal([]byte(value), &payload); err != nil {
		return jsonString(map[string]any{"value": value})
	}
	return jsonString(payload)
}

func automationFailureMetadataJSON(message string, metadata map[string]any) string {
	payload := make(map[string]any, len(metadata)+1)
	for key, value := range metadata {
		key = strings.TrimSpace(key)
		if key == "" || value == nil {
			continue
		}
		payload[key] = value
	}
	payload["message"] = strings.TrimSpace(message)
	return jsonString(payload)
}

func jsonString(value any) string { bytes, _ := json.Marshal(value); return string(bytes) }