package main

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"github.com/gofiber/fiber/v2"
)

var TogetherErrorCodes map[string]string

func init() {
	TogetherErrorCodes = make(map[string]string)
	TogetherErrorCodes["400"] = "Invalid Request - Please contact the support."
	TogetherErrorCodes["401"] = "Invalid Authentication - Ensure that the API key is still valid."
	TogetherErrorCodes["403"] = "Set max_tokens to a lower number. Or leave it empty for using max value."
	TogetherErrorCodes["404"] = "Model not found."
	TogetherErrorCodes["429"] = "Rate limit reached for requests - You are sending requests too quickly."
	TogetherErrorCodes["500"] = "Issue on Together AI servers - Retry your request after a brief wait and contact Together AI if the issue persists."
	TogetherErrorCodes["503"] = "Servers are experiencing high traffic - Please retry your requests after a brief wait."
	TogetherErrorCodes["504"] = "Servers are experiencing high traffic - Please retry your requests after a brief wait."
	TogetherErrorCodes["520"] = "An unexpected error has occurred internal to Together’s systems."
	TogetherErrorCodes["524"] = "An unexpected error has occurred internal to Together’s systems."
	TogetherErrorCodes["529"] = "An unexpected error has occurred internal to Together’s systems."
}

type TogetherChatCompletionResponse struct {
	ID      string           `json:"id"`
	Object  string           `json:"object"`
	Created int64            `json:"created"`
	Model   string           `json:"model"`
	Usage   OpenaiUsage      `json:"usage"`
	Choices []TogetherChoice `json:"choices"`
}

type TogetherChoice struct {
	Text         string `json:"text"`
	FinishReason string `json:"finish_reason"`
	Index        int    `json:"index"`
}

func TestTogetherKey(apiKey string) bool {
	url := "https://api.together.xyz/v1/completions"

	// Convert messages to OpenAI format
	togetherMessages := []RequestMessage{
		{
			Role:    "user",
			Content: "Hello",
		},
	}

	requestBody := OpenaiChatCompletionRequest{
		Model:       "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
		Messages:    togetherMessages,
		Temperature: 0,
		MaxTokens:   10,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return false
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return false
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return false
	}

	var chatCompletionResponse OpenaiChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return false
	}
	if chatCompletionResponse.Usage.CompletionTokens == 0 {
		return false
	}
	return true
}

func RequestTogether(c *fiber.Ctx, llm LLM, messages []Message) string {
	model := llm.Model.ModelID
	temperature := float64(llm.Temperature)
	context := llm.Context
	maxTokens := int(llm.MaxToken)

	var apiKey string
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		with
			filtered_keys := (
				select Key {
					key
				} filter .company.name = <str>$0 AND .<keys[is Setting].<setting[is User] = global currentUser
			)
		select filtered_keys.key limit 1
	`, &apiKey, "together")
	if err != nil {
		return "JADE internal error: 07-00-0000. Please contact the support."
	}

	url := "https://api.together.xyz/v1/completions"

	requestBody := OpenaiChatCompletionRequest{
		Model:       model,
		Messages:    Message2RequestMessage(messages, context),
		MaxTokens:   maxTokens,
		Temperature: temperature,
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return "JADE internal error: 07-01-0001. Please contact the support."
	}

	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonBody))
	if err != nil {
		return "JADE internal error: 07-02-0002. Please contact the support."
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return "JADE internal error: 07-02-0003. Please contact the support."
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "JADE internal error: 07-01-0004. Please contact the support."
	}

	for key, value := range TogetherErrorCodes {
		if strings.Contains(resp.Status, key) {
			return value
		}
	}

	var chatCompletionResponse TogetherChatCompletionResponse
	err = json.Unmarshal(body, &chatCompletionResponse)
	if err != nil {
		return "JADE internal error: 07-01-0005. Please contact the support."
	}

	var usedModelInfo ModelInfo
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		SELECT ModelInfo {
			inputPrice,
			outputPrice
		}
		FILTER .modelID = <str>$0
		LIMIT 1
	`, &usedModelInfo, model)
	if err != nil {
		return "JADE internal error: 07-00-0006. Please contact the support."
	}

	var inputCost float32 = float32(chatCompletionResponse.Usage.PromptTokens) * usedModelInfo.InputPrice
	var outputCost float32 = float32(chatCompletionResponse.Usage.CompletionTokens) * usedModelInfo.OutputPrice
	addUsage(c, inputCost, outputCost, chatCompletionResponse.Usage.PromptTokens, chatCompletionResponse.Usage.CompletionTokens, model)

	if len(chatCompletionResponse.Choices) == 0 {
		logMissingErrorCode.Println("Together -", resp.Status, "-", string(body))
		return "JADE internal error: 07-03-0007. Please contact the support."
	}

	return chatCompletionResponse.Choices[0].Text
}