// Package llmjson tolerates the small JSON artifacts that LLMs sometimes emit
// (e.g. `...` placeholders, trailing or duplicated commas) so structured
// outputs can still be parsed instead of failing the whole job.
package llmjson

import (
	"encoding/json"
	"strings"
)

// Unmarshal parses payload into v, retrying once on a sanitized copy when the
// raw bytes are not strictly valid JSON. The original error is returned when
// the sanitized retry also fails, so callers keep a meaningful message.
func Unmarshal(payload []byte, v any) error {
	err := json.Unmarshal(payload, v)
	if err == nil {
		return nil
	}
	if cleaned := Sanitize(payload); cleaned != nil {
		if retryErr := json.Unmarshal(cleaned, v); retryErr == nil {
			return nil
		}
	}
	return err
}

// Sanitize removes ellipsis placeholders, escapes unescaped control characters
// inside JSON string literals, and normalizes stray commas that sit outside of
// string literals. It returns nil when nothing changed so callers can avoid a
// redundant re-parse. String meaning is preserved; only invalid JSON escapes
// are added.
func Sanitize(in []byte) []byte {
	normalized := normalizeCommas(stripEllipsis(escapeControlCharsInStrings(string(in))))
	if normalized == string(in) {
		return nil
	}
	return []byte(normalized)
}

func escapeControlCharsInStrings(s string) string {
	var b strings.Builder
	b.Grow(len(s) + 16)
	inString := false
	escaped := false
	for i := 0; i < len(s); i++ {
		ch := s[i]
		if inString {
			if escaped {
				escaped = false
				b.WriteByte(ch)
				continue
			}
			switch ch {
			case '\\':
				escaped = true
				b.WriteByte(ch)
			case '"':
				inString = false
				b.WriteByte(ch)
			case '\n':
				b.WriteString(`\n`)
			case '\r':
				b.WriteString(`\r`)
			case '\t':
				b.WriteString(`\t`)
			default:
				if ch < 0x20 {
					b.WriteString(`\u00`)
					b.WriteByte(hexDigit(ch >> 4))
					b.WriteByte(hexDigit(ch & 0x0f))
				} else {
					b.WriteByte(ch)
				}
			}
			continue
		}
		if ch == '"' {
			inString = true
		}
		b.WriteByte(ch)
	}
	return b.String()
}

func hexDigit(v byte) byte {
	if v < 10 {
		return '0' + v
	}
	return 'a' + (v - 10)
}

func stripEllipsis(s string) string {
	var b strings.Builder
	b.Grow(len(s))
	inString := false
	escaped := false
	for i := 0; i < len(s); i++ {
		ch := s[i]
		if inString {
			b.WriteByte(ch)
			if escaped {
				escaped = false
			} else if ch == '\\' {
				escaped = true
			} else if ch == '"' {
				inString = false
			}
			continue
		}
		if ch == '"' {
			inString = true
			b.WriteByte(ch)
			continue
		}
		if ch == '.' {
			j := i
			for j < len(s) && s[j] == '.' {
				j++
			}
			if j-i >= 2 {
				i = j - 1
				continue
			}
		}
		b.WriteByte(ch)
	}
	return b.String()
}

func normalizeCommas(s string) string {
	var b strings.Builder
	b.Grow(len(s))
	inString := false
	escaped := false
	var lastMeaningful byte
	for i := 0; i < len(s); i++ {
		ch := s[i]
		if inString {
			b.WriteByte(ch)
			if escaped {
				escaped = false
			} else if ch == '\\' {
				escaped = true
			} else if ch == '"' {
				inString = false
			}
			continue
		}
		if ch == '"' {
			inString = true
			b.WriteByte(ch)
			lastMeaningful = ch
			continue
		}
		if ch == ',' {
			if lastMeaningful == '[' || lastMeaningful == '{' || lastMeaningful == ',' || lastMeaningful == 0 {
				continue
			}
			if next := nextNonSpaceByte(s, i+1); next == ']' || next == '}' {
				continue
			}
			b.WriteByte(ch)
			lastMeaningful = ch
			continue
		}
		b.WriteByte(ch)
		if ch != ' ' && ch != '\t' && ch != '\n' && ch != '\r' {
			lastMeaningful = ch
		}
	}
	return b.String()
}

func nextNonSpaceByte(s string, from int) byte {
	for i := from; i < len(s); i++ {
		switch s[i] {
		case ' ', '\t', '\n', '\r':
			continue
		default:
			return s[i]
		}
	}
	return 0
}