package mcp

import (
	"bytes"
	"encoding/json"
	"math"
	"strconv"
	"strings"
)

var kbForbiddenInputFields = makeStringSet(
	"target", "target_key", "thread_channel_id", "message_id",
	"kb_ids", "kbIds", "source_ids", "sourceIds", "source_id", "sourceId",
	"repo_id", "repoId", "repo_role", "repoRole", "path", "ref", "revision",
	"route", "citation_uri", "citationUri", "search_mode", "searchMode", "offset",
	"title_only", "titleOnly", "slug_prefix", "slugPrefix", "items", "metadata_only",
)

type tolerantInt struct {
	value int
}

func (value *tolerantInt) UnmarshalJSON(data []byte) error {
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.UseNumber()
	var raw any
	if err := decoder.Decode(&raw); err != nil {
		return nil
	}
	number, ok := raw.(json.Number)
	if !ok {
		return nil
	}
	parsed, err := strconv.ParseFloat(number.String(), 64)
	if err != nil || math.IsNaN(parsed) || math.IsInf(parsed, 0) {
		return nil
	}
	value.value = int(parsed)
	return nil
}

type tolerantStrings struct {
	values []string
}

func (value *tolerantStrings) UnmarshalJSON(data []byte) error {
	var raw []json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return nil
	}
	seen := make(map[string]struct{}, len(raw))
	value.values = nil
	for _, item := range raw {
		var text string
		if err := json.Unmarshal(item, &text); err != nil {
			continue
		}
		text = strings.TrimSpace(text)
		if text == "" {
			continue
		}
		if _, duplicate := seen[text]; duplicate {
			continue
		}
		seen[text] = struct{}{}
		value.values = append(value.values, text)
	}
	return nil
}

type kbSearchInput struct {
	Query      string
	MaxResults tolerantInt
	extra      map[string]json.RawMessage
}

func (input *kbSearchInput) UnmarshalJSON(data []byte) error {
	fields, err := decodeKBInputFields(data)
	if err != nil {
		return err
	}
	_ = json.Unmarshal(fields["query"], &input.Query)
	_ = json.Unmarshal(fields["max_results"], &input.MaxResults)
	delete(fields, "query")
	delete(fields, "max_results")
	input.extra = fields
	return nil
}

func (input kbSearchInput) normalized() (kbSearchInput, error) {
	input.Query = strings.TrimSpace(input.Query)
	if input.Query == "" {
		return kbSearchInput{}, stringError("query is required")
	}
	input.MaxResults.value = normalizeKBAgentTeamMaxResults(input.MaxResults.value)
	return input, nil
}

func (input kbSearchInput) MarshalJSON() ([]byte, error) {
	fields := cloneKBSafeExtras(input.extra)
	fields["query"] = mustJSONRaw(input.Query)
	fields["max_results"] = mustJSONRaw(input.MaxResults.value)
	fields["limit"] = mustJSONRaw(input.MaxResults.value)
	return json.Marshal(fields)
}

type kbFetchInput struct {
	CitationIDs     tolerantStrings
	MaxCharsPerItem tolerantInt
	extra           map[string]json.RawMessage
}

func (input *kbFetchInput) UnmarshalJSON(data []byte) error {
	fields, err := decodeKBInputFields(data)
	if err != nil {
		return err
	}
	_ = json.Unmarshal(fields["citation_ids"], &input.CitationIDs)
	_ = json.Unmarshal(fields["max_chars_per_item"], &input.MaxCharsPerItem)
	delete(fields, "citation_ids")
	delete(fields, "max_chars_per_item")
	input.extra = fields
	return nil
}

func (input kbFetchInput) normalized() (kbFetchInput, error) {
	if len(input.CitationIDs.values) == 0 {
		return kbFetchInput{}, stringError("citation_ids is required")
	}
	return input, nil
}

func (input kbFetchInput) MarshalJSON() ([]byte, error) {
	fields := cloneKBSafeExtras(input.extra)
	fields["citation_ids"] = mustJSONRaw(input.CitationIDs.values)
	fields["citationIds"] = mustJSONRaw(input.CitationIDs.values)
	if input.MaxCharsPerItem.value > 0 {
		fields["max_chars_per_item"] = mustJSONRaw(input.MaxCharsPerItem.value)
		fields["maxBytes"] = mustJSONRaw(input.MaxCharsPerItem.value)
	}
	return json.Marshal(fields)
}

func decodeKBInputFields(data []byte) (map[string]json.RawMessage, error) {
	if len(data) == 0 || string(data) == "null" {
		data = []byte(`{}`)
	}
	fields := make(map[string]json.RawMessage)
	if err := json.Unmarshal(data, &fields); err != nil {
		return nil, stringError("invalid tool arguments")
	}
	for forbidden := range kbForbiddenInputFields {
		delete(fields, forbidden)
	}
	return fields, nil
}

func cloneKBSafeExtras(values map[string]json.RawMessage) map[string]json.RawMessage {
	cloned := make(map[string]json.RawMessage, len(values)+4)
	for key, value := range values {
		if _, forbidden := kbForbiddenInputFields[key]; forbidden {
			continue
		}
		cloned[key] = append(json.RawMessage(nil), value...)
	}
	return cloned
}

func mustJSONRaw(value any) json.RawMessage {
	encoded, err := json.Marshal(value)
	if err != nil {
		panic(err)
	}
	return encoded
}

func makeStringSet(values ...string) map[string]struct{} {
	result := make(map[string]struct{}, len(values))
	for _, value := range values {
		result[value] = struct{}{}
	}
	return result
}