package main

import (
	"fmt"
	"regexp"
	"strings"
)

var exitCodeMarkerRe = regexp.MustCompile(`^\[(\d+|\*)\]$`)

func rewriteFile(file ShoutFile, results []CommandResult, originalContent string) string {
	lines := strings.Split(originalContent, "\n")
	var output []string

	cmdIdx := 0

	for i := 0; i < len(lines); i++ {
		line := lines[i]

		if strings.HasPrefix(line, "$ ") {
			output = append(output, line)

			if cmdIdx >= len(file.Commands) || cmdIdx >= len(results) {
				// Out of bounds — skip past expected output, preserving original lines
				j := i + 1
				for j < len(lines) && !strings.HasPrefix(lines[j], "$ ") {
					j++
				}
				output = append(output, lines[i+1:j]...)
				i = j - 1
				cmdIdx++
				continue
			}

			cmd := file.Commands[cmdIdx]
			result := results[cmdIdx]

			// Skip past old expected output lines
			j := i + 1
			for j < len(lines) && !strings.HasPrefix(lines[j], "$ ") {
				j++
			}

			oldExpectedRaw := lines[i+1 : j]

			// Count trailing blank lines
			trailingBlanks := 0
			for k := len(oldExpectedRaw) - 1; k >= 0; k-- {
				if oldExpectedRaw[k] == "" {
					trailingBlanks++
				} else {
					break
				}
			}

			outputMatch := matchOutput(cmd.Expected, result.Actual)
			ecMatch := exitCodeOK(cmd.ExitCodeType, cmd.ExitCodeValue, result.ExitCode)

			if outputMatch && ecMatch {
				// Both match — preserve original verbatim
				output = append(output, oldExpectedRaw...)
			} else if outputMatch {
				// Output matches but exit code changed — preserve output, fix exit code
				raw := oldExpectedRaw
				rawTrailing := 0
				for k := len(raw) - 1; k >= 0; k-- {
					if raw[k] == "" {
						rawTrailing++
					} else {
						break
					}
				}
				content := raw[:len(raw)-rawTrailing]
				// Strip old exit code marker if present
				if len(content) > 0 && exitCodeMarkerRe.MatchString(content[len(content)-1]) {
					content = content[:len(content)-1]
				}
				output = append(output, content...)
				marker := updateExitCodeMarker(cmd, result.ExitCode)
				if marker != "" {
					output = append(output, marker)
				} else if len(content) > 0 && exitCodeMarkerRe.MatchString(content[len(content)-1]) {
					output = append(output, "[0]")
				}
				for k := 0; k < rawTrailing; k++ {
					output = append(output, "")
				}
			} else {
				// Output doesn't match — rewrite everything
				output = append(output, result.Actual...)
				marker := updateExitCodeMarker(cmd, result.ExitCode)
				if marker != "" {
					output = append(output, marker)
				} else if len(result.Actual) > 0 && exitCodeMarkerRe.MatchString(result.Actual[len(result.Actual)-1]) {
					output = append(output, "[0]")
				}
				for k := 0; k < trailingBlanks; k++ {
					output = append(output, "")
				}
			}

			i = j - 1
			cmdIdx++
		} else {
			output = append(output, line)
		}
	}

	return strings.Join(output, "\n")
}

// updateExitCodeMarker returns the exit code marker to write, preserving [*]
// wildcards when the command still exits non-zero.
func updateExitCodeMarker(cmd Command, actualExitCode int) string {
	if cmd.ExitCodeType == ExitCodeWildcard && actualExitCode != 0 {
		return "[*]"
	}
	if actualExitCode != 0 {
		return fmt.Sprintf("[%d]", actualExitCode)
	}
	if cmd.ExitCodeType == ExitCodeWildcard || cmd.ExitCodeType == ExitCodeExact {
		return "[0]"
	}
	return ""
}