package llm

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/pkoukk/tiktoken-go"
	"github.com/tidbcloud/AgentTeam/internal/apr/tokenencoding"
)

type countTokensRequest struct {
	Model    string               `json:"model"`
	Messages []countTokensMessage `json:"messages,omitempty"`
	System   any                  `json:"system,omitempty"`
	Tools    []countTokensTool    `json:"tools,omitempty"`
}

type countTokensMessage struct {
	Role    string `json:"role"`
	Content any    `json:"content"`
}

type countTokensTool struct {
	Name        string         `json:"name"`
	Description *string        `json:"description,omitempty"`
	InputSchema map[string]any `json:"input_schema,omitempty"`
}

type countTokensResponse struct {
	InputTokens int `json:"input_tokens"`
}

type countTokenizer struct {
	encoder   *tiktoken.Tiktoken
	constants countModelConstants
}

type countModelConstants struct {
	FuncInit int
	PropInit int
	PropKey  int
	EnumInit int
	EnumItem int
	FuncEnd  int
}

type countTokensHandler struct {
	mu    sync.RWMutex
	cache map[string]*countTokenizer
}

var defaultCountTokensHandler = &countTokensHandler{cache: make(map[string]*countTokenizer)}

func (handler *LLMHandler) CountTokens(ctx *gin.Context) {
	resolved, ok := handler.resolveCountTokensRequest(ctx)
	if !ok {
		return
	}
	_, release, ok := handler.acquireConcurrency(ctx, resolved, "count_tokens")
	if !ok {
		return
	}
	defer release()
	defaultCountTokensHandler.handleResolved(ctx, resolved)
}

func (handler *countTokensHandler) handleResolved(ctx *gin.Context, resolved *requestResolver) {
	var req countTokensRequest
	if err := json.Unmarshal(resolved.Body, &req); err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid request body: %v", err)})
		return
	}
	if strings.TrimSpace(req.Model) == "" {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "model field is required"})
		return
	}
	tk, err := handler.getTokenizer(req.Model)
	if err != nil {
		ctx.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to get tokenizer: %v", err)})
		return
	}
	count, err := tk.CountTokens(&req)
	if err != nil {
		ctx.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to count tokens: %v", err)})
		return
	}
	if len(req.Tools) > 0 {
		anthropicBeta := ctx.GetHeader("anthropic-beta")
		mcpToolExist := false
		if anthropicBeta != "" && strings.HasPrefix(anthropicBeta, "claude-code") {
			for _, tool := range req.Tools {
				if strings.HasPrefix(tool.Name, "mcp__") {
					mcpToolExist = true
					break
				}
			}
		}
		if !mcpToolExist {
			if strings.HasPrefix(req.Model, "claude") {
				count += 346
			} else if strings.HasPrefix(req.Model, "grok") {
				count += 480
			}
		}
	}
	if strings.HasPrefix(req.Model, "claude") {
		count = int(float64(count) * 1.15)
	} else if strings.HasPrefix(req.Model, "grok") {
		count = int(float64(count) * 1.03)
	}
	if count < 1 {
		count = 1
	}
	ctx.JSON(http.StatusOK, countTokensResponse{InputTokens: count})
}

func (handler *LLMHandler) resolveCountTokensRequest(ctx *gin.Context) (*requestResolver, bool) {
	rawBody, err := io.ReadAll(ctx.Request.Body)
	if err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "failed to read request body"})
		return nil, false
	}
	var request struct {
		Model string `json:"model"`
	}
	if err := json.Unmarshal(rawBody, &request); err != nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
		return nil, false
	}
	authorized, ok := llmProxyAuthorizationFromContext(ctx.Request.Context())
	if !ok {
		ctx.JSON(http.StatusUnauthorized, gin.H{"error": "missing llm proxy authorization"})
		return nil, false
	}
	provider, normalizedModel := handler.resolveAuthorizedProviderModel(authorized, request.Model)
	if provider == nil {
		ctx.JSON(http.StatusBadRequest, gin.H{"error": "unsupported model: " + request.Model})
		return nil, false
	}
	rewrittenBody, err := rewriteModelField(rawBody, normalizedModel)
	if err != nil {
		ctx.JSON(http.StatusInternalServerError, gin.H{"error": "failed to rewrite request model"})
		return nil, false
	}
	return &requestResolver{Provider: provider, Model: normalizedModel, Body: rewrittenBody}, true
}

func (handler *countTokensHandler) getTokenizer(model string) (*countTokenizer, error) {
	handler.mu.RLock()
	if tk, ok := handler.cache[model]; ok {
		handler.mu.RUnlock()
		return tk, nil
	}
	handler.mu.RUnlock()
	handler.mu.Lock()
	defer handler.mu.Unlock()
	if tk, ok := handler.cache[model]; ok {
		return tk, nil
	}
	tk, err := newCountTokenizer(model)
	if err != nil {
		return nil, err
	}
	handler.cache[model] = tk
	return tk, nil
}

func newCountTokenizer(model string) (*countTokenizer, error) {
	encoding, constants := getCountModelEncoding(model)
	encoder, err := tokenencoding.GetEncoding(encoding)
	if err != nil {
		encoder, err = tokenencoding.GetEncoding("o200k_base")
		if err != nil {
			return nil, fmt.Errorf("failed to get encoding: %w", err)
		}
	}
	return &countTokenizer{encoder: encoder, constants: constants}, nil
}

func getCountModelEncoding(model string) (string, countModelConstants) {
	switch {
	case strings.Contains(model, "gpt"):
		return "o200k_base", countModelConstants{FuncInit: 10, PropInit: 3, PropKey: 3, EnumInit: -3, EnumItem: 3, FuncEnd: 12}
	case strings.Contains(model, "claude"):
		return "cl100k_base", countModelConstants{FuncInit: 7, PropInit: 3, PropKey: 3, EnumInit: -3, EnumItem: 3, FuncEnd: 12}
	default:
		return "o200k_base", countModelConstants{FuncInit: 7, PropInit: 3, PropKey: 3, EnumInit: -3, EnumItem: 3, FuncEnd: 12}
	}
}

func (tokenizer *countTokenizer) CountTokens(req *countTokensRequest) (int, error) {
	total := 0
	if req.System != nil {
		total += tokenizer.countSystem(req.System)
	}
	for _, message := range req.Messages {
		total += tokenizer.countMessage(message)
	}
	if len(req.Tools) > 0 {
		total += tokenizer.countTools(req.Tools)
	}
	total += 3
	return total, nil
}

func (tokenizer *countTokenizer) countSystem(system any) int {
	switch value := system.(type) {
	case string:
		return len(tokenizer.encoder.Encode(value, nil, nil))
	case []any:
		tokens := 0
		for _, item := range value {
			if block, ok := item.(map[string]any); ok {
				if text, ok := block["text"].(string); ok {
					tokens += len(tokenizer.encoder.Encode(text, nil, nil))
				}
			}
		}
		return tokens
	default:
		return 0
	}
}

func (tokenizer *countTokenizer) countMessage(message countTokensMessage) int {
	tokens := 3
	tokens += len(tokenizer.encoder.Encode(message.Role, nil, nil))
	tokens += tokenizer.countContent(message.Content)
	return tokens
}

func (tokenizer *countTokenizer) countContent(content any) int {
	switch value := content.(type) {
	case string:
		return len(tokenizer.encoder.Encode(value, nil, nil))
	case []any:
		tokens := 0
		for _, item := range value {
			block, ok := item.(map[string]any)
			if !ok {
				continue
			}
			tokens += tokenizer.countContentBlock(block)
		}
		return tokens
	default:
		return 0
	}
}

func (tokenizer *countTokenizer) countContentBlock(block map[string]any) int {
	typ, _ := block["type"].(string)
	switch typ {
	case "text":
		text, _ := block["text"].(string)
		return len(tokenizer.encoder.Encode(text, nil, nil))
	case "image", "image_url":
		placeholder := "image"
		if imageURL, ok := block["image_url"].(map[string]any); ok {
			if url, ok := imageURL["url"].(string); ok && url != "" {
				placeholder = url
			}
		}
		return len(tokenizer.encoder.Encode(placeholder, nil, nil)) + 85
	case "tool_use":
		name, _ := block["name"].(string)
		id, _ := block["id"].(string)
		inputJSON, _ := json.Marshal(block["input"])
		return len(tokenizer.encoder.Encode(name, nil, nil)) + len(tokenizer.encoder.Encode(id, nil, nil)) + len(tokenizer.encoder.Encode(string(inputJSON), nil, nil))
	case "tool_result":
		id, _ := block["tool_use_id"].(string)
		return len(tokenizer.encoder.Encode(id, nil, nil)) + tokenizer.countContent(block["content"])
	case "thinking":
		thinking, _ := block["thinking"].(string)
		return len(tokenizer.encoder.Encode(thinking, nil, nil))
	default:
		blob, _ := json.Marshal(block)
		return len(tokenizer.encoder.Encode(string(blob), nil, nil))
	}
}

func (tokenizer *countTokenizer) countTools(tools []countTokensTool) int {
	if len(tools) == 0 {
		return 0
	}
	tokens := tokenizer.constants.FuncInit
	for _, tool := range tools {
		tokens += tokenizer.countSingleTool(tool)
	}
	tokens += tokenizer.constants.FuncEnd
	return tokens
}

func (tokenizer *countTokenizer) countSingleTool(tool countTokensTool) int {
	tokens := tokenizer.constants.FuncInit
	desc := ""
	if tool.Description != nil {
		desc = strings.TrimSuffix(*tool.Description, ".")
	}
	line := tool.Name + ":" + desc
	tokens += len(tokenizer.encoder.Encode(line, nil, nil))
	if tool.InputSchema != nil {
		tokens += tokenizer.countParameters(tool.InputSchema)
	}
	return tokens
}

func (tokenizer *countTokenizer) countParameters(params map[string]any) int {
	tokens := 0
	if properties, ok := params["properties"].(map[string]any); ok && len(properties) > 0 {
		tokens += tokenizer.constants.PropInit
		for key, value := range properties {
			tokens += tokenizer.countParameter(key, value)
		}
	}
	for key, value := range params {
		if key == "properties" {
			continue
		}
		paramBytes, _ := json.Marshal(value)
		tokens += len(tokenizer.encoder.Encode(fmt.Sprintf("%s:%s", key, string(paramBytes)), nil, nil))
	}
	return tokens
}

func (tokenizer *countTokenizer) countParameter(key string, param any) int {
	tokens := tokenizer.constants.PropKey
	paramObj, ok := param.(map[string]any)
	if !ok {
		return tokens
	}
	paramType, _ := paramObj["type"].(string)
	if paramType == "" {
		paramType = "string"
	}
	paramDesc, _ := paramObj["description"].(string)
	paramDesc = strings.TrimSuffix(paramDesc, ".")
	if enumValues, ok := paramObj["enum"].([]any); ok {
		tokens += tokenizer.constants.EnumInit
		for _, item := range enumValues {
			tokens += tokenizer.constants.EnumItem
			tokens += len(tokenizer.encoder.Encode(fmt.Sprintf("%v", item), nil, nil))
		}
	}
	line := fmt.Sprintf("%s:%s:%s", key, paramType, paramDesc)
	tokens += len(tokenizer.encoder.Encode(line, nil, nil))
	for name, value := range paramObj {
		if name == "type" || name == "description" || name == "enum" {
			continue
		}
		paramBytes, _ := json.Marshal(value)
		tokens += len(tokenizer.encoder.Encode(fmt.Sprintf("%s:%s", name, string(paramBytes)), nil, nil))
	}
	return tokens
}