package util

import (
	"io"
	"regexp"

	"github.com/fatih/color"
	yaml "github.com/oasdiff/yaml3"
)

// YAMLOutput writes data as YAML to the given writer.
// If pretty is true, output is indented with 2 spaces.
// If colorize is true, YAML syntax is colorized.
func YAMLOutput(w io.Writer, data any, pretty, colorize bool) error {
	encoder := yaml.NewEncoder(w)
	defer encoder.Close()

	if pretty {
		encoder.SetIndent(2)
	}

	if !colorize {
		return encoder.Encode(data)
	}

	// For colored output, encode to bytes first, then colorize
	var buf []byte
	bufWriter := &byteWriter{buf: &buf}
	bufEncoder := yaml.NewEncoder(bufWriter)
	if pretty {
		bufEncoder.SetIndent(2)
	}
	if err := bufEncoder.Encode(data); err != nil {
		return err
	}
	bufEncoder.Close()

	colorized := colorizeYAML(string(buf))
	_, err := w.Write([]byte(colorized))
	return err
}

type byteWriter struct {
	buf *[]byte
}

func (b *byteWriter) Write(p []byte) (n int, err error) {
	*b.buf = append(*b.buf, p...)
	return len(p), nil
}

// colorizeYAML applies ANSI colors to YAML output.
func colorizeYAML(input string) string {
	// Color definitions - wrap to match ReplaceAllStringFunc signature
	keyColor := func(s string) string { return color.New(color.FgCyan).Sprint(s) }
	stringColor := func(s string) string { return color.New(color.FgGreen).Sprint(s) }
	numberColor := func(s string) string { return color.New(color.FgYellow).Sprint(s) }
	boolColor := func(s string) string { return color.New(color.FgMagenta).Sprint(s) }
	nullColor := func(s string) string { return color.New(color.FgRed).Sprint(s) }

	// Patterns for YAML elements
	keyPattern := regexp.MustCompile(`(?m)^(\s*)([a-zA-Z_][a-zA-Z0-9_-]*)(:)`)
	stringPattern := regexp.MustCompile(`:\s*"([^"]*)"`)
	quotedStringPattern := regexp.MustCompile(`:\s*'([^']*)'`)
	numberPattern := regexp.MustCompile(`:\s*(-?\d+\.?\d*)\s*$`)
	boolPattern := regexp.MustCompile(`:\s*(true|false)\s*$`)
	nullPattern := regexp.MustCompile(`:\s*(null|~)\s*$`)

	result := input

	// Apply colors in order (specific patterns first)
	result = nullPattern.ReplaceAllStringFunc(result, func(s string) string {
		return regexp.MustCompile(`(null|~)`).ReplaceAllStringFunc(s, nullColor)
	})

	result = boolPattern.ReplaceAllStringFunc(result, func(s string) string {
		return regexp.MustCompile(`(true|false)`).ReplaceAllStringFunc(s, boolColor)
	})

	result = numberPattern.ReplaceAllStringFunc(result, func(s string) string {
		return regexp.MustCompile(`(-?\d+\.?\d*)\s*$`).ReplaceAllStringFunc(s, numberColor)
	})

	result = stringPattern.ReplaceAllStringFunc(result, func(s string) string {
		return regexp.MustCompile(`"([^"]*)"`).ReplaceAllStringFunc(s, stringColor)
	})

	result = quotedStringPattern.ReplaceAllStringFunc(result, func(s string) string {
		return regexp.MustCompile(`'([^']*)'`).ReplaceAllStringFunc(s, stringColor)
	})

	result = keyPattern.ReplaceAllStringFunc(result, func(s string) string {
		matches := keyPattern.FindStringSubmatch(s)
		if len(matches) >= 4 {
			return matches[1] + keyColor(matches[2]) + matches[3]
		}
		return s
	})

	return result
}