// 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 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/URL contents are never altered.
func Sanitize(in []byte) []byte {
	normalized := normalizeCommas(stripEllipsis(string(in)))
	if normalized == string(in) {
		return nil
	}
	return []byte(normalized)
}

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
}