package server

import (
	"encoding/json"
	"fmt"

	"github.com/mark3labs/mcp-go/mcp"
	"go.uber.org/zap"

	"github.com/smart-mcp-proxy/mcpproxy-go/internal/truncate"
)

// CacheStore is the narrow interface forwardContentResult needs from a cache
// manager to persist truncated payloads so the read_cache MCP tool can later
// page through them. It mirrors cache.Manager.Store but is defined here to
// keep this package free of an internal/cache import.
type CacheStore interface {
	Store(key, toolName string, args map[string]interface{}, content, recordPath string, totalRecords int) error
}

// logCacheStoreFailure emits a structured warning when a CacheStore.Store call
// fails. The truncation banner shown to the agent will reference a cache key
// that no longer resolves, so the symptom (read_cache returning "cache key not
// found") needs to be debuggable from the server logs. A nil logger is a no-op
// so test callers can pass nil without ceremony.
func logCacheStoreFailure(logger *zap.Logger, err error, toolName, cacheKey string, totalRecords, payloadBytes int) {
	if logger == nil {
		return
	}
	logger.Warn("Failed to persist truncated payload — read_cache will return 'cache key not found' for this key",
		zap.String("tool", toolName),
		zap.String("cache_key", cacheKey),
		zap.Int("total_records", totalRecords),
		zap.Int("payload_bytes", payloadBytes),
		zap.Error(err),
	)
}

// forwardContentResult preserves non-text content blocks (ImageContent, AudioContent,
// EmbeddedResource) from an upstream CallToolResult while applying truncation only to
// TextContent blocks. This fixes issue #368 where all content types were being
// serialized to JSON and wrapped in a single TextContent, destroying the ability of
// vision-capable LLMs to process images efficiently.
//
// When truncation kicks in, the FULL pre-truncation content is persisted via
// cacheStore (if non-nil) under the same key embedded in the truncated text's
// "use read_cache" instruction. Without this step read_cache always returns
// "cache key not found", which was the long-standing pagination bug.
//
// Parameters:
//   - result: the upstream CallToolResult (passed as interface{} from Manager.CallTool)
//   - truncator: applies size limits to text blocks only
//   - cacheStore: optional; when non-nil, full pre-truncation payloads are stored here
//     under the cache key generated by the truncator so read_cache can paginate them
//   - logger: optional; receives a zap.Warn entry if a cacheStore.Store call fails.
//     A failed Store leaves the agent with a banner pointing at a key that won't
//     resolve, so this log is the only way the operator can diagnose the resulting
//     "cache key not found" symptom. Pass nil to silence.
//   - toolName, args: passed to truncator for cache-key derivation and to the cacheStore
//
// Returns:
//   - forwarded: the CallToolResult to return downstream (with original content types)
//   - textRepresentation: a text rendering of the full result for logging/metrics
//   - wasTruncated: whether any TextContent block was truncated
//
// If result is not a *mcp.CallToolResult, it falls back to JSON-serializing the whole
// thing into a TextContent block (legacy behavior).
func forwardContentResult(result interface{}, truncator *truncate.Truncator, cacheStore CacheStore, logger *zap.Logger, toolName string, args map[string]interface{}) (forwarded *mcp.CallToolResult, textRepresentation string, wasTruncated bool) {
	ctr, ok := result.(*mcp.CallToolResult)
	if !ok || ctr == nil {
		// Fallback: not a CallToolResult (should not happen with current upstream chain,
		// but guard against future changes). Use legacy JSON-wrap behavior.
		jsonBytes, err := json.Marshal(result)
		if err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("Failed to serialize result: %v", err)), "", false
		}
		text := string(jsonBytes)
		if truncator != nil && truncator.ShouldTruncate(text) {
			tr := truncator.Truncate(text, toolName, args)
			if tr.CacheAvailable && cacheStore != nil {
				if storeErr := cacheStore.Store(tr.CacheKey, toolName, args, text, tr.RecordPath, tr.TotalRecords); storeErr != nil {
					logCacheStoreFailure(logger, storeErr, toolName, tr.CacheKey, tr.TotalRecords, len(text))
				}
			}
			text = tr.TruncatedContent
			wasTruncated = true
		}
		return mcp.NewToolResultText(text), text, wasTruncated
	}

	// Walk the Content slice, truncating only TextContent blocks.
	// Build a parallel text representation for logging/metrics.
	newContent := make([]mcp.Content, 0, len(ctr.Content))
	var textBuilder []string
	for _, c := range ctr.Content {
		switch tc := c.(type) {
		case mcp.TextContent:
			txt := tc.Text
			if truncator != nil && truncator.ShouldTruncate(txt) {
				tr := truncator.Truncate(txt, toolName, args)
				if tr.CacheAvailable && cacheStore != nil {
					if storeErr := cacheStore.Store(tr.CacheKey, toolName, args, txt, tr.RecordPath, tr.TotalRecords); storeErr != nil {
						logCacheStoreFailure(logger, storeErr, toolName, tr.CacheKey, tr.TotalRecords, len(txt))
					}
				}
				txt = tr.TruncatedContent
				wasTruncated = true
			}
			tc.Text = txt
			newContent = append(newContent, tc)
			textBuilder = append(textBuilder, txt)
		case mcp.ImageContent:
			// Preserve image blocks unchanged. For logging, emit a placeholder.
			newContent = append(newContent, tc)
			textBuilder = append(textBuilder, fmt.Sprintf("[image:%s len=%d]", tc.MIMEType, len(tc.Data)))
		case mcp.AudioContent:
			// Preserve audio blocks unchanged. For logging, emit a placeholder.
			newContent = append(newContent, tc)
			textBuilder = append(textBuilder, fmt.Sprintf("[audio:%s len=%d]", tc.MIMEType, len(tc.Data)))
		default:
			// Unknown content type (e.g., EmbeddedResource). Forward as-is and
			// best-effort JSON encode for logging.
			newContent = append(newContent, c)
			if b, err := json.Marshal(c); err == nil {
				textBuilder = append(textBuilder, string(b))
			}
		}
	}

	forwarded = &mcp.CallToolResult{
		Result:            ctr.Result,
		Content:           newContent,
		StructuredContent: ctr.StructuredContent,
		IsError:           ctr.IsError,
	}
	textRepresentation = joinTextParts(textBuilder)
	return forwarded, textRepresentation, wasTruncated
}

// maybeTruncateAndCacheText is a small helper for built-in MCP tools whose
// output bypasses forwardContentResult (e.g. handleReadCache). It applies the
// same truncate-and-cache contract: if the text exceeds the truncator's limit,
// the full text is stored in cacheStore under the truncator-generated key and
// the returned snippet carries the standard "use read_cache" banner that
// points at it.
//
// The paginableUnits parameter is the recursion guard for built-in tools that
// may themselves be paginating (notably read_cache). If paginableUnits <= 1
// the caller has nothing further to subdivide — caching a fresh key here would
// just hand the agent a new key resolving to the same oversize payload, an
// inescapable loop. In that case the text is returned as-is (still over the
// limit) and the caller decides what to do.
//
// Termination is therefore agent-driven, not strictly algorithmic: a level
// whose records all fit individually but overflow collectively keeps minting
// a fresh key per call until the agent pages down to the single-record
// short-circuit. Each call's key is unique (timestamped at nanosecond
// granularity), so successive identical calls are NOT idempotent — they
// produce distinct cache entries. Unbounded growth is bounded only by the
// cache's TTL + background cleanup (see cache.DefaultTTL), which is acceptable
// because truncation is the exceptional path.
//
// Returns:
//   - the text to actually emit (possibly truncated)
//   - wasTruncated: true if the truncator's recorded-pagination path ran
//
// logger is optional; receives a zap.Warn if cacheStore.Store fails so the
// resulting "cache key not found" can be debugged. Pass nil to silence.
func maybeTruncateAndCacheText(text, toolName string, args map[string]interface{}, paginableUnits int, truncator *truncate.Truncator, cacheStore CacheStore, logger *zap.Logger) (out string, wasTruncated bool) {
	if truncator == nil || !truncator.ShouldTruncate(text) {
		return text, false
	}
	if paginableUnits <= 1 {
		// Cannot subdivide further; recursive caching here would just hand the
		// agent a new key that resolves to the exact same payload. Return the
		// oversize text intact and let the caller (and the agent) handle it.
		return text, false
	}
	tr := truncator.Truncate(text, toolName, args)
	if tr.CacheAvailable && cacheStore != nil {
		if storeErr := cacheStore.Store(tr.CacheKey, toolName, args, text, tr.RecordPath, tr.TotalRecords); storeErr != nil {
			logCacheStoreFailure(logger, storeErr, toolName, tr.CacheKey, tr.TotalRecords, len(text))
		}
	}
	return tr.TruncatedContent, true
}

// forwardedText recomputes the text representation of a (possibly sanitised)
// forwarded result so logging/metrics reflect what the agent actually receives
// after Spec 054 Track B mutation. It mirrors forwardContentResult's text
// rendering (text verbatim, placeholders for non-text blocks). If forwarded is
// nil or carries no renderable content, the supplied fallback is returned.
func forwardedText(forwarded *mcp.CallToolResult, fallback string) string {
	if forwarded == nil {
		return fallback
	}
	var parts []string
	for _, c := range forwarded.Content {
		switch tc := c.(type) {
		case mcp.TextContent:
			parts = append(parts, tc.Text)
		case mcp.ImageContent:
			parts = append(parts, fmt.Sprintf("[image:%s len=%d]", tc.MIMEType, len(tc.Data)))
		case mcp.AudioContent:
			parts = append(parts, fmt.Sprintf("[audio:%s len=%d]", tc.MIMEType, len(tc.Data)))
		default:
			if b, err := json.Marshal(c); err == nil {
				parts = append(parts, string(b))
			}
		}
	}
	if len(parts) == 0 {
		return fallback
	}
	return joinTextParts(parts)
}

// joinTextParts concatenates text parts with a newline separator.
// Equivalent to strings.Join but avoids importing strings just for this.
func joinTextParts(parts []string) string {
	switch len(parts) {
	case 0:
		return ""
	case 1:
		return parts[0]
	}
	total := len(parts) - 1 // separators
	for _, p := range parts {
		total += len(p)
	}
	out := make([]byte, 0, total)
	for i, p := range parts {
		if i > 0 {
			out = append(out, '\n')
		}
		out = append(out, p...)
	}
	return string(out)
}
