package tool

import (
	"bytes"
	"encoding/json"
	"fmt"
	"regexp"
)

// fencedJSONRe matches a code-fence wrapper some models emit around tool-call
// JSON, e.g. ```json {...} ```. The captured group is the inner JSON.
var fencedJSONRe = regexp.MustCompile("(?s)^[\\s`]*```(?:json)?[\\r\\n]*(.*?)[\\r\\n]*```\\s*$")

// LenientJSONUnmarshal unmarshals raw into dest, applying a few tolerant repairs
// for the malformed tool-call JSON models sometimes emit, before falling back to
// strict encoding/json. Repairs are applied in order:
//
//  1. strip a wrapping ```json code fence (and stray backticks);
//  2. trim surrounding whitespace;
//  3. strip trailing commas before `}` and `]`;
//  4. attempt the strict decode.
//
// The function reports the raw json.Unmarshal error (wrapped) when every repair
// fails, so a caller can fall back to treating the payload as plain text without
// crashing — the malformed call never aborts the round. Full repair (unquoted
// keys, single-quoted strings, smart quotes) is deliberately left for a later
// milestone; the load-bearing guarantee here is that a malformed envelope is
// recoverable or degrades to text rather than throwing.
func LenientJSONUnmarshal(raw []byte, dest any) error {
	candidates := repairCandidates(raw)
	var lastErr error
	for _, c := range candidates {
		if err := json.Unmarshal(c, dest); err == nil {
			return nil
		} else {
			lastErr = err
		}
	}
	return fmt.Errorf("tool: lenient json: %w", lastErr)
}

// repairCandidates returns the ordered list of byte slices to try, deduplicated:
// the raw input, the fenced-stripped input, and trailing-comma-stripped variants.
func repairCandidates(raw []byte) [][]byte {
	out := [][]byte{nil}
	seen := map[string]bool{}
	push := func(b []byte) {
		s := string(b)
		if seen[s] {
			return
		}
		seen[s] = true
		out = append(out, b)
	}

	push(bytes.TrimSpace(raw))

	clean := raw
	if m := fencedJSONRe.FindSubmatch(raw); len(m) == 2 {
		clean = bytes.TrimSpace(m[1])
		push(clean)
	}
	clean = bytes.TrimSpace(clean)
	push(bytes.TrimSpace(clean))
	push(bytes.TrimSpace(stripTrailingCommas(clean)))
	push(bytes.TrimSpace(stripTrailingCommas(bytes.TrimSpace(raw))))

	return out
}

// stripTrailingCommas removes commas that sit immediately before a closing `}`
// or `]` (ignoring whitespace), the most common single-token repair for truncated
// model output such as `{"a":1,}`.
func stripTrailingCommas(b []byte) []byte {
	if !bytes.ContainsAny(b, "}]") {
		return b
	}
	// Scan for a comma run whose next non-whitespace byte is a closing bracket; a
	// copy is built so the original bytes are never mutated.
	out := make([]byte, 0, len(b))
	skipComma := false
	for i := 0; i < len(b); i++ {
		if b[i] == ',' {
			// Peek ahead past whitespace for a closing bracket.
			j := i + 1
			for j < len(b) && isJSONSpace(b[j]) {
				j++
			}
			if j < len(b) && (b[j] == '}' || b[j] == ']') {
				skipComma = true
			}
		}
		if skipComma && isJSONSpace(b[i]) {
			continue
		}
		if skipComma && b[i] == ',' {
			skipComma = false
			continue
		}
		out = append(out, b[i])
	}
	return out
}

// isJSONSpace reports whether b is JSON insignificant whitespace.
func isJSONSpace(b byte) bool {
	switch b {
	case ' ', '\t', '\r', '\n':
		return true
	default:
		return false
	}
}