package llm

import (
	"encoding/json"
	"log/slog"
	"strings"

	anthropic "github.com/anthropics/anthropic-sdk-go"
	"github.com/firebase/genkit/go/ai"
	openai "github.com/openai/openai-go"
	openaiParam "github.com/openai/openai-go/packages/param"
	"github.com/openai/openai-go/responses"
)

const (
	removedToolResultPlaceholder   = "Tool result removed to fit context."
	contextWindowExceededError     = "context exceeds model window after removing tool results"
	contextOutputReserveTokenCount = 8192
	defaultContextWindowTokenCount = 200000
)

type contextReduction struct {
	OriginalTokenCount int
	ReducedTokenCount  int
	RemovedToolResults int
	FitsBudget         bool
}

type toolResultReductionTarget struct {
	tokenCount int
	remove     func()
}

func estimateContextTokenCount(text string) int {
	if text == "" {
		return 0
	}
	return max(1, (len(text)+2)/3)
}

func contextInputBudget(contextWindowTokenCount int) int {
	window := contextWindowTokenCount
	if window <= 0 {
		window = defaultContextWindowTokenCount
	}
	safety := max(1024, window/100)
	budget := window - contextOutputReserveTokenCount - safety
	if budget < 1 {
		return 1
	}
	return budget
}

func contextFitsBudget(contextWindowTokenCount int, currentInputTokenCount int) bool {
	return currentInputTokenCount <= contextInputBudget(contextWindowTokenCount)
}

func reduceToolResultsForRequest(contextWindowTokenCount int, currentInputTokenCount int, targets []toolResultReductionTarget) contextReduction {
	budget := contextInputBudget(contextWindowTokenCount)
	reduction := newContextReduction(currentInputTokenCount, budget)

	placeholderTokenCount := estimateContextTokenCount(removedToolResultPlaceholder)
	for _, target := range targets {
		if reduction.ReducedTokenCount <= budget {
			break
		}
		target.remove()
		reduction.ReducedTokenCount -= target.tokenCount
		reduction.ReducedTokenCount += placeholderTokenCount
		reduction.RemovedToolResults++
	}

	reduction.FitsBudget = reduction.ReducedTokenCount <= budget
	slog.Debug(
		"Reduced context tool results",
		"inputTokenCount", reduction.OriginalTokenCount,
		"reducedTokenCount", reduction.ReducedTokenCount,
		"budgetTokenCount", budget,
		"removedToolResultCount", reduction.RemovedToolResults,
		"fitsBudget", reduction.FitsBudget,
	)
	return reduction
}

func newContextReduction(currentInputTokenCount int, budget int) contextReduction {
	return contextReduction{
		OriginalTokenCount: currentInputTokenCount,
		ReducedTokenCount:  currentInputTokenCount,
		FitsBudget:         currentInputTokenCount <= budget,
	}
}

func estimateJSONTokenCount(v any) int {
	b, err := json.Marshal(v)
	if err != nil {
		return 0
	}
	return estimateContextTokenCount(string(b))
}

func reduceOpenAIContextForRequest(contextWindowTokenCount int,
	currentInputTokenCount int,
	next []openai.ChatCompletionMessageParamUnion,
) ([]openai.ChatCompletionMessageParamUnion, contextReduction) {
	targets := make([]toolResultReductionTarget, 0)
	for i, msg := range next {
		if msg.OfTool == nil {
			continue
		}
		content := openAIToolContent(msg.OfTool.Content)
		if content == removedToolResultPlaceholder {
			continue
		}
		idx := i
		targets = append(targets, toolResultReductionTarget{
			tokenCount: estimateContextTokenCount(content),
			remove: func() {
				next[idx].OfTool.Content.OfString = openaiParam.NewOpt(removedToolResultPlaceholder)
				next[idx].OfTool.Content.OfArrayOfContentParts = nil
			},
		})
	}

	return next, reduceToolResultsForRequest(contextWindowTokenCount, currentInputTokenCount, targets)
}

func openAIToolContent(content openai.ChatCompletionToolMessageParamContentUnion) string {
	if content.OfString.Valid() {
		return content.OfString.Value
	}
	return string(marshalContextOrEmpty(content))
}

func reduceResponsesContextForRequest(contextWindowTokenCount int,
	currentInputTokenCount int,
	next []responses.ResponseInputItemUnionParam,
) ([]responses.ResponseInputItemUnionParam, contextReduction) {
	targets := make([]toolResultReductionTarget, 0)
	for i, item := range next {
		if item.OfFunctionCallOutput == nil || item.OfFunctionCallOutput.Output == removedToolResultPlaceholder {
			continue
		}
		idx := i
		targets = append(targets, toolResultReductionTarget{
			tokenCount: estimateContextTokenCount(item.OfFunctionCallOutput.Output),
			remove: func() {
				next[idx].OfFunctionCallOutput.Output = removedToolResultPlaceholder
			},
		})
	}

	return next, reduceToolResultsForRequest(contextWindowTokenCount, currentInputTokenCount, targets)
}

func reduceAnthropicContextForRequest(contextWindowTokenCount int,
	currentInputTokenCount int,
	next []anthropic.MessageParam,
) ([]anthropic.MessageParam, contextReduction) {
	targets := make([]toolResultReductionTarget, 0)
	for mi := range next {
		for bi, block := range next[mi].Content {
			if block.OfToolResult == nil {
				continue
			}
			content := anthropicToolResultContent(block.OfToolResult)
			if content == removedToolResultPlaceholder {
				continue
			}
			messageIdx := mi
			blockIdx := bi
			targets = append(targets, toolResultReductionTarget{
				tokenCount: estimateContextTokenCount(content),
				remove: func() {
					next[messageIdx].Content[blockIdx].OfToolResult.Content = []anthropic.ToolResultBlockParamContentUnion{
						{
							OfText: &anthropic.TextBlockParam{
								Text: removedToolResultPlaceholder,
							},
						},
					}
				},
			})
		}
	}

	return next, reduceToolResultsForRequest(contextWindowTokenCount, currentInputTokenCount, targets)
}

func anthropicToolResultContent(result *anthropic.ToolResultBlockParam) string {
	var b strings.Builder
	for _, content := range result.Content {
		if content.OfText != nil {
			b.WriteString(content.OfText.Text)
			continue
		}
		b.Write(marshalContextOrEmpty(content))
	}
	return b.String()
}

func reduceGenkitContextForRequest(contextWindowTokenCount int,
	currentInputTokenCount int,
	next []*ai.Message,
) ([]*ai.Message, contextReduction) {
	targets := make([]toolResultReductionTarget, 0)
	for mi, msg := range next {
		if msg == nil || msg.Role != ai.RoleTool {
			continue
		}
		for pi, part := range msg.Content {
			if part == nil || part.ToolResponse == nil || part.ToolResponse.Output == removedToolResultPlaceholder {
				continue
			}
			messageIdx := mi
			partIdx := pi
			targets = append(targets, toolResultReductionTarget{
				tokenCount: estimateJSONTokenCount(part.ToolResponse.Output),
				remove: func() {
					next[messageIdx].Content[partIdx].ToolResponse.Output = removedToolResultPlaceholder
				},
			})
		}
	}

	return next, reduceToolResultsForRequest(contextWindowTokenCount, currentInputTokenCount, targets)
}

func marshalContextOrEmpty(v any) []byte {
	b, err := json.Marshal(v)
	if err != nil {
		return nil
	}
	return b
}
