package bot

import (
	"encoding/json"

	"queryorchestration/internal/serviceconfig/chatbot"
)

// StripCachedResults removes the top-level "cached_results" key from a
// session-state JSON object and returns the re-marshaled bytes. Other
// top-level keys are preserved verbatim. Plan §7: "Non-null
// updated_session_state strips cached_results before persisting."
//
// The function is intentionally non-recursive: only the top-level
// cached_results key is removed. Aaria's session_state schema places
// cached_results at the top level (see plans/chatbot.stuff/sample.response.json),
// so a recursive strip would be both unnecessary and risky — nothing
// inside resolved_entities or last_query_plan should be modified.
//
// Edge cases:
//   - nil/empty input round-trips as nil so the caller can pass NULL
//     state without a special case.
//   - JSON null input round-trips as JSON null.
//   - Input that is not a JSON object returns the original bytes
//     unchanged plus a non-nil error so callers can decide whether to
//     log + persist nil or fail closed.
func StripCachedResults(state json.RawMessage) (json.RawMessage, error) {
	if len(state) == 0 {
		return nil, nil
	}
	// JSON null handled explicitly so the caller does not see a 4-byte
	// "null" string round-tripped through map decode + remarshal.
	if string(state) == "null" {
		return state, nil
	}

	var m map[string]any
	if err := json.Unmarshal(state, &m); err != nil {
		return state, err
	}
	delete(m, "cached_results")
	return json.Marshal(m)
}

// SessionStateOversize reports whether the supplied state byte length
// exceeds chatbot.StateMaxBytes. Plan §8 caps persisted session state
// at 64 KiB. The caller is responsible for invoking StripCachedResults
// first so the cap reflects the post-strip footprint.
func SessionStateOversize(state json.RawMessage) bool {
	return len(state) > chatbot.StateMaxBytes
}