package knowledgeqa

import (
	"encoding/json"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

const genericGroundingDiagnosticCode = "knowledge_qa.grounding.warning"

func projectRuntimeActivityPayload(activity *protocol.RuntimeActivity, raw []byte) {
	if activity == nil {
		return
	}
	var payload map[string]json.RawMessage
	if err := model.UnmarshalJSONObject("Knowledge Q&A event payload", raw, &payload); err != nil {
		return
	}

	// Preserve the legacy model-call projection for older clients while the
	// structured stage projection rolls out.
	if activity.ToolName == "model" {
		if modelCalls := nonNegativePayloadInt(payload, "modelCalls"); modelCalls != nil {
			activity.ModelCalls = *modelCalls
		}
	}

	stage, ok := payloadString(payload, "stage")
	if !ok {
		return
	}
	switch stage {
	case "scope_authorized":
		counts := &protocol.RuntimeActivityCounts{
			AuthorizedKnowledgeBases:  nonNegativePayloadInt(payload, "authorizedKnowledgeBaseCount"),
			AuthorizedSources:         nonNegativePayloadInt(payload, "authorizedSourceCount"),
			MissingKnowledgeBases:     nonNegativePayloadInt(payload, "missingKnowledgeBases"),
			MissingSources:            nonNegativePayloadInt(payload, "missingSources"),
			UnavailableKnowledgeBases: nonNegativePayloadInt(payload, "unavailableKnowledgeBases"),
			UnavailableSources:        nonNegativePayloadInt(payload, "unavailableSources"),
			PartiallyAvailableSources: nonNegativePayloadInt(payload, "partiallyAvailableSources"),
		}
		activity.Stage = stage
		activity.Counts = nonEmptyRuntimeActivityCounts(counts)
	case "history_read":
		counts := &protocol.RuntimeActivityCounts{
			ContextTurns: nonNegativePayloadInt(payload, "tailTurnCount"),
		}
		activity.Stage = stage
		activity.Counts = nonEmptyRuntimeActivityCounts(counts)
	case "exploration_progress":
		counts := &protocol.RuntimeActivityCounts{
			ModelCalls: nonNegativePayloadInt(payload, "modelCalls"),
		}
		activity.Stage = stage
		activity.Counts = nonEmptyRuntimeActivityCounts(counts)
	case "retrieval_complete":
		counts := &protocol.RuntimeActivityCounts{
			SearchCalls:    nonNegativePayloadInt(payload, "grepCalls"),
			ListCalls:      nonNegativePayloadInt(payload, "listCalls"),
			ReadCalls:      nonNegativePayloadInt(payload, "readCalls"),
			WebSearchCalls: nonNegativePayloadInt(payload, "webSearchCalls"),
			WebReadCalls:   nonNegativePayloadInt(payload, "webReadCalls"),
			WebSources:     nonNegativePayloadInt(payload, "webSources"),
			RefCalls:       nonNegativePayloadInt(payload, "refCalls"),
			ModelCalls:     nonNegativePayloadInt(payload, "modelCalls"),
			SearchedRoots:  nonNegativePayloadInt(payload, "searchedRoots"),
		}
		activity.Stage = stage
		activity.Counts = nonEmptyRuntimeActivityCounts(counts)
	case "grounding_diagnostic":
		activity.Stage = stage
		activity.DiagnosticCode = projectedGroundingDiagnosticCode(payload)
	}
}

func payloadString(payload map[string]json.RawMessage, key string) (string, bool) {
	var value string
	if err := json.Unmarshal(payload[key], &value); err != nil {
		return "", false
	}
	value = strings.TrimSpace(value)
	return value, value != ""
}

func nonNegativePayloadInt(payload map[string]json.RawMessage, key string) *int {
	var value int
	if err := json.Unmarshal(payload[key], &value); err != nil || value < 0 {
		return nil
	}
	return &value
}

func projectedGroundingDiagnosticCode(payload map[string]json.RawMessage) string {
	code, _ := payloadString(payload, "code")
	switch code {
	case "knowledge_qa.grounding.claim_empty",
		"knowledge_qa.grounding.claim_evidence_missing",
		"knowledge_qa.grounding.evidence_unavailable",
		"knowledge_qa.grounding.marker_mismatch",
		genericGroundingDiagnosticCode:
		return code
	default:
		return genericGroundingDiagnosticCode
	}
}

func nonEmptyRuntimeActivityCounts(counts *protocol.RuntimeActivityCounts) *protocol.RuntimeActivityCounts {
	if counts == nil ||
		(counts.AuthorizedKnowledgeBases == nil &&
			counts.AuthorizedSources == nil &&
			counts.MissingKnowledgeBases == nil &&
			counts.MissingSources == nil &&
			counts.UnavailableKnowledgeBases == nil &&
			counts.UnavailableSources == nil &&
			counts.PartiallyAvailableSources == nil &&
			counts.ContextTurns == nil &&
			counts.SearchCalls == nil &&
			counts.ListCalls == nil &&
			counts.ReadCalls == nil &&
			counts.WebSearchCalls == nil &&
			counts.WebReadCalls == nil &&
			counts.WebSources == nil &&
			counts.RefCalls == nil &&
			counts.ModelCalls == nil &&
			counts.SearchedRoots == nil) {
		return nil
	}
	return counts
}