// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Package lifecycle implements image install progress reporting.
package lifecycle

import (
	"fmt"
	"maps"
	"slices"
	"sort"
	"strings"
	"unicode/utf8"

	"github.com/siderolabs/talos/pkg/machinery/api/machine"
	"github.com/siderolabs/talos/pkg/machinery/constants"
	"github.com/siderolabs/talos/pkg/reporter"
)

// ProgressWriter writes install progress updates to a reporter.
type ProgressWriter struct {
	// ongoingInstalls keeps track of ongoing install jobs per node.
	ongoingInstalls map[string]installJob
	captureOutput   bool
	output          map[string]*outputBuffer
}

const maxCapturedOutputSize = 64 * 1024

// NewProgressWriter initializes a progress writer, optionally retaining a bounded tail of installer output.
func NewProgressWriter(captureOutput bool) *ProgressWriter {
	return &ProgressWriter{captureOutput: captureOutput}
}

// UpdateJob updates the progress of a pull job for a given node and layer ID.
//
// It is supposed to be called whenever there is a progress update for a layer pull.
func (w *ProgressWriter) UpdateJob(node string, status *machine.LifecycleServiceInstallProgress) {
	if w.ongoingInstalls == nil {
		w.ongoingInstalls = make(map[string]installJob)
	}

	if message, ok := status.GetResponse().(*machine.LifecycleServiceInstallProgress_Message); ok && w.captureOutput {
		if w.output == nil {
			w.output = make(map[string]*outputBuffer)
		}

		if w.output[node] == nil {
			w.output[node] = &outputBuffer{}
		}

		w.output[node].Append(message.Message)
	}

	w.ongoingInstalls[node] = installJob{Status: status}
}

func (w *ProgressWriter) outputForNode(node string) string {
	if w.output[node] == nil {
		return ""
	}

	return w.output[node].String()
}

// Failure formats the classified exit code for a failed operation and includes captured installer output when enabled.
func (w *ProgressWriter) Failure(node, operation string, exitCode int32) string {
	var result strings.Builder

	if output := strings.TrimSpace(w.outputForNode(node)); output != "" {
		fmt.Fprintf(&result, "%s: installer output:\n%s\n", node, output)
	}

	fmt.Fprintf(&result, "%s: %s failed with exit code %d", node, operation, exitCode)

	if description := exitCodeDescription(exitCode); description != "" {
		fmt.Fprintf(&result, " (%s)", description)
	}

	return result.String()
}

func exitCodeDescription(exitCode int32) string {
	switch int(exitCode) {
	case constants.ExitInvalidInput:
		return "invalid input"
	case constants.ExitUnsupported:
		return "unsupported operation"
	case constants.ExitEnvironment:
		return "environment error"
	case constants.ExitDependency:
		return "dependency error"
	case constants.ExitIO:
		return "I/O error"
	case constants.ExitInstall:
		return "installation error"
	default:
		return ""
	}
}

type outputBuffer struct {
	messages  []string
	size      int
	truncated bool
}

func (buffer *outputBuffer) Append(message string) {
	if message == "" {
		return
	}

	if len(message) > maxCapturedOutputSize {
		message = message[len(message)-maxCapturedOutputSize:]
		for len(message) > 0 && !utf8.RuneStart(message[0]) {
			message = message[1:]
		}

		message = strings.Clone(message)

		buffer.messages = nil
		buffer.size = 0
		buffer.truncated = true
	}

	for buffer.size+len(message) > maxCapturedOutputSize {
		buffer.size -= len(buffer.messages[0])
		buffer.messages[0] = ""
		buffer.messages = buffer.messages[1:]
		buffer.truncated = true
	}

	buffer.messages = append(buffer.messages, message)
	buffer.size += len(message)
}

func (buffer *outputBuffer) String() string {
	var result strings.Builder

	if buffer.truncated {
		result.WriteString("[earlier installer output truncated]\n")
	}

	for _, message := range buffer.messages {
		result.WriteString(message)
	}

	return result.String()
}

// PrintLayerProgress prints the current layer pull progress to the reporter.
func (w *ProgressWriter) PrintLayerProgress(rep *reporter.Reporter) {
	nodes := slices.Collect(maps.Keys(w.ongoingInstalls))
	sort.Strings(nodes)

	sb := strings.Builder{}

	for _, node := range nodes {
		sb.WriteString(node + ": ")

		fmt.Fprintf(&sb, "%s\n", w.ongoingInstalls[node].Status.Fmt())
	}

	rep.Report(reporter.Update{
		Message: sb.String(),
		Status:  reporter.StatusRunning,
	})
}

type installJob struct {
	Status *machine.LifecycleServiceInstallProgress
}
