package ucca

import "context"

// EvidenceType classifies a retrieved unit by WHAT KIND of evidence it is, independent of its
// collection. Footnotes/tables/figures are Evidence Types, not collections. The Authority Router
// surfaces non-text evidence from the authoritative knowledge space (the KB slice) SEPARATELY from
// the merged text top-K, so fine-grained evidence isn't outranked by broad-base text.
//
// The layer this introduces: Intent -> Knowledge Space -> EvidenceType -> Collection -> Merge ->
// Authority. Today FOOTNOTE is populated; FIGURE arrives with C8 and TABLE is already present from
// C6/C9 — no router rebuild needed, the same path carries every new evidence type.
type EvidenceType string

const (
	EvidenceText     EvidenceType = "text"
	EvidenceFootnote EvidenceType = "footnote"
	EvidenceTable    EvidenceType = "table"
	EvidenceFigure   EvidenceType = "figure"
)

// classifyEvidence derives the EvidenceType from a result's payload markers. Precedence
// footnote > figure > table > text (a unit carries at most one is_* marker in practice).
func classifyEvidence(r LegalSearchResult) EvidenceType {
	switch {
	case r.IsFootnote:
		return EvidenceFootnote
	case r.IsFigure:
		return EvidenceFigure
	case r.IsTable:
		return EvidenceTable
	default:
		return EvidenceText
	}
}

// evidenceRetrievalTopK is the budget for the authoritative-KB evidence pass. Deliberately targeted
// (the authoritative slice within the recognized knowledge space), NOT a blanket top-K increase of
// the merged result set — the successes came from BETTER-targeted evidence, not MORE evidence.
const evidenceRetrievalTopK = 20

// maxEvidencePerType caps each surfaced evidence type.
const maxEvidencePerType = 6

// RetrieveEvidence returns the authoritative typed evidence (footnotes/tables/figures) for an
// in-scope query, pulled from the KB slice and grouped by EvidenceType. This is the "Evidence Type"
// router layer (Option A): when the query is in the KB knowledge space, the authoritative evidence
// within that space is surfaced separately so it isn't lost in the broad-base text merge. Returns an
// empty map when out of scope or KB routing is disabled. Text evidence is NOT returned here — it
// flows through the normal Retrieve() merge (the LLM context + the sources list).
func (c *LegalRAGClient) RetrieveEvidence(ctx context.Context, query string) map[EvidenceType][]LegalSearchResult {
	ev := map[EvidenceType][]LegalSearchResult{}
	if !c.kbScopeRoutingEnabled || c.kbSliceCollection == "" || !inKBScope(query) {
		return ev
	}
	hits, err := c.searchInternal(ctx, c.kbSliceCollection, query, nil, evidenceRetrievalTopK)
	if err != nil {
		return ev
	}
	for _, h := range hits {
		t := classifyEvidence(h)
		if t == EvidenceText || len(ev[t]) >= maxEvidencePerType {
			continue
		}
		ev[t] = append(ev[t], h)
	}
	return ev
}
