package controller

import (
	"fmt"
	"net/http"
	"opencatd-open/internal/dto"
	"opencatd-open/internal/model"
	"opencatd-open/llm"
	"opencatd-open/llm/claude/v2"
	"opencatd-open/llm/google/v2"
	"opencatd-open/llm/openai_compatible"
	"opencatd-open/pkg/tokenizer"
	"strconv"

	"github.com/gin-gonic/gin"
)

func (h *Proxy) ChatHandler(c *gin.Context) {
	var chatreq llm.ChatRequest
	if err := c.ShouldBindJSON(&chatreq); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	err := h.SelectApiKey(chatreq.Model)
	if err != nil {
		dto.WrapErrorAsOpenAI(c, 500, err.Error())
		return
	}

	var llm llm.LLM
	switch *h.apikey.ApiType {
	case "claude":
		llm, err = claude.NewClaude(h.apikey)
	case "gemini":
		llm, err = google.NewGemini(c, h.apikey)
	case "openai", "azure", "github":
		fallthrough
	default:
		llm, err = openai_compatible.NewOpenAICompatible(h.apikey)
		if err != nil {
			dto.WrapErrorAsOpenAI(c, 500, fmt.Errorf("create llm client error: %w", err).Error())
			return
		}
	}

	if !chatreq.Stream {
		resp, err := llm.Chat(c, chatreq)
		if err != nil {
			dto.WrapErrorAsOpenAI(c, 500, err.Error())
		}
		c.JSON(http.StatusOK, resp)

	} else {
		datachan, err := llm.StreamChat(c, chatreq)
		if err != nil {
			dto.WrapErrorAsOpenAI(c, 500, err.Error())
		}
		for data := range datachan {
			c.SSEvent("", data)
		}
	}

	llmusage := llm.GetTokenUsage()

	cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens+llmusage.ToolsTokens, llmusage.CompletionTokens)
	userid, _ := strconv.ParseInt(c.GetString("user_id"), 10, 64)
	usage := model.Usage{
		UserID:           userid,
		Model:            llmusage.Model,
		Stream:           chatreq.Stream,
		PromptTokens:     llmusage.PromptTokens + llmusage.ToolsTokens,
		CompletionTokens: llmusage.CompletionTokens,
		TotalTokens:      llmusage.TotalTokens,
		Cost:             fmt.Sprintf("%f", cost),
	}
	h.SendUsage(&usage)
	defer fmt.Println("cost:", cost, "prompt_tokens:", llmusage.PromptTokens, "completion_tokens:", llmusage.CompletionTokens, "total_tokens:", llmusage.TotalTokens)

}