package tools

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/krow/krow-backend/go-api/internal/knowledge"
)

// knowledge_search: retrieval the model can drive.
//
// The loop already retrieves once, for the caller's opening question, and puts
// the result in a <context> block. That covers the common case and covers it
// cheaply — but it is one shot at one phrasing, and the phrasing is the user's.
// A question like "am I allowed to leave early on Fridays?" retrieves a
// paragraph about early departure and misses the one about shift-swap approvals
// that actually answers it.
//
// So the model gets a second bite: it may search again, in its own words, once
// it knows what it is looking for. That is worth a tool.
//
// What it emphatically does NOT get is a widening of scope. The sources come
// from the agent's spec by way of ctx.KnowledgeSources, exactly as the
// principal does; there is no `source` argument, because an argument is
// something a model can choose and the set of corpora an agent may read is not
// the model's to choose. All the model controls is the words.

type knowledgeSearchInput struct {
	Query string `json:"query"`
	Limit int    `json:"limit"`
}

// KnowledgeSearch builds the search tool.
//
// One registered tool, not one per agent: the per-agent part is the source
// list, and that travels on the Context beside the principal rather than being
// closed over. Both are set by the loop from records the conversation cannot
// touch, so two agents sharing this tool still cannot read each other's
// corpora — and the registry stays a flat set of names, which is what §3's
// publish-time validation resolves against.
func KnowledgeSearch(r *knowledge.Retriever) Tool {
	return Tool{
		Name: "knowledge_search",
		Description: "Search the documents this agent has access to and get back passages " +
			"with source ids. Use it when the answer depends on what a written policy, " +
			"handbook or guide actually says — and search again in your own words if the " +
			"first passages are close but not quite right. Cite the source id of anything " +
			"you rely on.",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"query": map[string]any{
					"type": "string",
					"description": "What to look for. Write it as the words you would expect to " +
						"find in the document, not as a question to a person.",
				},
				"limit": map[string]any{
					"type": "integer", "minimum": 1, "maximum": 20,
					"description": "How many passages to return. Defaults to 8.",
				},
			},
			"required":             []string{"query"},
			"additionalProperties": false,
		},
		Effect:         EffectRead,
		MaxResultBytes: DefaultMaxResultBytes,
		Handler: func(ctx context.Context, tc Context, inputs json.RawMessage) Result {
			if r == nil {
				return Failf(CodeUnavailable, "there are no documents to search")
			}
			if len(tc.KnowledgeSources) == 0 {
				// This agent's spec named no knowledge. Refused rather than
				// widened: an empty source list is not permission to read
				// everything, and the tool should not have been offered.
				return Failf(CodeUnavailable, "this agent has no documents to search")
			}

			var in knowledgeSearchInput
			if err := json.Unmarshal(inputs, &in); err != nil {
				return Failf(CodeInvalidInput, "the arguments were not valid JSON")
			}
			if strings.TrimSpace(in.Query) == "" {
				return Failf(CodeInvalidInput, "a search needs something to search for")
			}

			limit := in.Limit
			if limit <= 0 {
				limit = knowledge.DefaultK
			}
			if limit > 20 {
				limit = 20
			}

			// The principal is the caller's. Not the agent's, not a service
			// account, and not anything the model supplied — I1 lives in this
			// one line as much as anywhere in the package.
			res, err := r.Retrieve(ctx, knowledge.Query{
				Text:      in.Query,
				Principal: tc.Principal,
				Sources:   tc.KnowledgeSources,
				K:         limit,
			})
			if err != nil {
				var kErr *knowledge.Error
				if ok := asKnowledge(err, &kErr); ok && kErr.Code == knowledge.ErrNoPrincipal {
					// A caller retrieval will not serve is refused the same way
					// every other resource refuses one. Saying "you have no
					// tenant" would be a more useful error and a worse one.
					return Denied()
				}
				return Failf(CodeFailed, "the documents could not be searched")
			}

			passages := make([]map[string]any, 0, len(res.Chunks))
			for _, c := range res.Chunks {
				p := map[string]any{
					// The id first, because citing it is the point. §5: a claim
					// without a retrievable citation is inference, not grounded
					// fact, and the model can only tell them apart if every
					// passage arrived with an address.
					"sourceId": c.ChunkID,
					"title":    c.Title,
					"text":     c.Text,
				}
				if c.Heading != "" {
					p["section"] = c.Heading
				}
				if c.URI != "" {
					p["uri"] = c.URI
				}
				passages = append(passages, p)
			}

			data := map[string]any{
				"query":    in.Query,
				"passages": passages,
				"count":    len(passages),
			}
			if len(passages) == 0 {
				data["note"] = "Nothing in these documents matched. This is a real answer, " +
					"not a failure to look — say so rather than answering from general knowledge."
			}
			if res.DenseSkipped != "" {
				// Handed to the model because it changes what an absence means.
				// "I found nothing" is a weaker claim when half the index was
				// not searched, and the model should be able to say which.
				data["degraded"] = res.DenseSkipped
			}
			return OK(data)
		},
	}
}

// asKnowledge is errors.As for the knowledge package's error type.
func asKnowledge(err error, target **knowledge.Error) bool {
	if e, ok := err.(*knowledge.Error); ok {
		*target = e
		return true
	}
	return false
}